fix(keys): fix range checks on input (#93)

Author: Joshua-Anderson

cmd/keys.go

@@ -2,7 +2,9 @@ package cmd
 
 import (
 	"fmt"
+	"io"
 	"io/ioutil"
+	"os"
 	"path/filepath"
 	"strconv"
 	"strings"
@@ -68,13 +70,19 @@ func KeyAdd(keyLocation string) error {
 	var key api.KeyCreateRequest
 
 	if keyLocation == "" {
-		key, err = chooseKey()
+		ks, err := listKeys()
+		if err != nil {
+			return err
+		}
+		key, err = chooseKey(ks, os.Stdin)
+		if err != nil {
+			return err
+		}
 	} else {
 		key, err = getKey(keyLocation)
-	}
-
-	if err != nil {
-		return err
+		if err != nil {
+			return err
+		}
 	}
 
 	fmt.Printf("Uploading %s to deis...", filepath.Base(key.Name))
@@ -88,13 +96,7 @@ func KeyAdd(keyLocation string) error {
 	return nil
 }
 
-func chooseKey() (api.KeyCreateRequest, error) {
-	keys, err := listKeys()
-
-	if err != nil {
-		return api.KeyCreateRequest{}, err
-	}
-
+func chooseKey(keys []api.KeyCreateRequest, input io.Reader) (api.KeyCreateRequest, error) {
 	fmt.Println("Found the following SSH public keys:")
 
 	for i, key := range keys {
@@ -106,15 +108,15 @@ func chooseKey() (api.KeyCreateRequest, error) {
 	var selected string
 
 	fmt.Print("Which would you like to use with Deis? ")
-	fmt.Scanln(&selected)
+	fmt.Fscanln(input, &selected)
 
 	numSelected, err := strconv.Atoi(selected)
 
 	if err != nil {
-		return api.KeyCreateRequest{}, err
+		return api.KeyCreateRequest{}, fmt.Errorf("%s is not a valid integer", selected)
 	}
 
-	if numSelected > len(keys)+1 {
+	if numSelected < 0 || numSelected > len(keys) {
 		return api.KeyCreateRequest{}, fmt.Errorf("%d is not a valid option", numSelected)
 	}
 

cmd/keys_test.go

@@ -1,10 +1,12 @@
 package cmd
 
 import (
+	"io"
 	"io/ioutil"
 	"os"
 	"path/filepath"
 	"reflect"
+	"strings"
 	"testing"
 
 	"github.com/deis/controller-sdk-go/api"
@@ -121,3 +123,37 @@ func TestListKeys(t *testing.T) {
 		t.Errorf("Expected %v, Got %v", expected, keys)
 	}
 }
+
+type chooseKeyCases struct {
+	Reader   io.Reader
+	Expected string
+}
+
+func TestChooseKey(t *testing.T) {
+	testKeys := []api.KeyCreateRequest{
+		api.KeyCreateRequest{},
+	}
+
+	checks := []chooseKeyCases{
+		chooseKeyCases{
+			Reader:   strings.NewReader("-1"),
+			Expected: "-1 is not a valid option",
+		},
+		chooseKeyCases{
+			Reader:   strings.NewReader("2"),
+			Expected: "2 is not a valid option",
+		},
+		chooseKeyCases{
+			Reader:   strings.NewReader("a"),
+			Expected: "a is not a valid integer",
+		},
+	}
+
+	for _, check := range checks {
+		_, err := chooseKey(testKeys, check.Reader)
+
+		if err.Error() != check.Expected {
+			t.Errorf("Expected %s, Got %s", check.Expected, err.Error())
+		}
+	}
+}