Added timeout testing and very basic get parameters testing (#47)

cmd/get.go

@@ -16,7 +16,7 @@ var format string
 
 func get(path string) {
 	p := util.NewParameterStorePath(path)
-	pms := io.ReadFromParameterStore(*p)
+	pms := io.ReadFromParameterStore(*p, nil)
 	if format == "yaml" || format == "yml" {
 		serialized, err := yaml.Marshal(pms)
 		if err != nil {

cmd/put.go

@@ -2,17 +2,25 @@ package cmd
 
 import (
 	"log"
+	"strconv"
+	"time"
 
 	"github.com/pbs/gorson/internal/gorson/io"
 	"github.com/pbs/gorson/internal/gorson/util"
 	"github.com/spf13/cobra"
 )
 
 var filename string
+var timeout string
 
-func put(path string, parameters map[string]string) {
+func put(path string, parameters map[string]string, timeout string) {
 	p := util.NewParameterStorePath(path)
-	err := io.WriteToParameterStore(parameters, *p)
+	timeoutInt, err := strconv.ParseInt(timeout, 0, 64)
+	timeoutDuration := time.Duration(timeoutInt) * time.Minute
+	if err != nil {
+		log.Fatal(err)
+	}
+	err = io.WriteToParameterStore(parameters, *p, timeoutDuration, nil)
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -25,11 +33,12 @@ func init() {
 		Run: func(cmd *cobra.Command, args []string) {
 			path := args[0]
 			parameters := io.ReadJSONFile(filename)
-			put(path, parameters)
+			put(path, parameters, timeout)
 		},
 		Args: cobra.ExactArgs(1),
 	}
 	cmd.Flags().StringVarP(&filename, "file", "f", "", "json file to read key/value pairs from")
+	cmd.Flags().StringVarP(&timeout, "timeout", "t", "1", "timeout in minutes for put")
 	cmd.MarkFlagRequired("file")
 	rootCmd.AddCommand(cmd)
 }

internal/gorson/io/io.go

@@ -30,8 +30,11 @@ func getSSMClient() *ssm.SSM {
 }
 
 // ReadFromParameterStore gets all parameters from a given parameter store path
-func ReadFromParameterStore(path util.ParameterStorePath) map[string]string {
-	client := getSSMClient()
+func ReadFromParameterStore(path util.ParameterStorePath, client ssmiface.SSMAPI) map[string]string {
+	if client == nil {
+		client = getSSMClient()
+	}
+
 	p := path.String()
 
 	var nextToken *string
@@ -120,8 +123,10 @@ func writeSingleParameter(c chan WriteResult, client ssmiface.SSMAPI, name strin
 }
 
 // WriteToParameterStore writes given parameters to a given parameter store path
-func WriteToParameterStore(parameters map[string]string, path util.ParameterStorePath) error {
-	client := getSSMClient()
+func WriteToParameterStore(parameters map[string]string, path util.ParameterStorePath, timeout time.Duration, client ssmiface.SSMAPI) error {
+	if client == nil {
+		client = getSSMClient()
+	}
 
 	// the jobs channel will receive messages from successful parameter store writes
 	jobs := make(chan WriteResult, len(parameters))
@@ -158,7 +163,7 @@ func WriteToParameterStore(parameters map[string]string, path util.ParameterStor
 			}
 		}
 		return nil
-	case <-time.After(1 * time.Minute):
+	case <-time.After(timeout):
 		return errors.New("timeout")
 	}
 }

internal/gorson/io/io_test.go

@@ -2,42 +2,101 @@ package io
 
 import (
 	"errors"
+	"reflect"
 	"testing"
+	"time"
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/awserr"
 	"github.com/aws/aws-sdk-go/service/ssm"
 	"github.com/aws/aws-sdk-go/service/ssm/ssmiface"
+	"github.com/pbs/gorson/internal/gorson/util"
 )
 
 type mockedPutParameterReturnPair struct {
 	Resp ssm.PutParameterOutput
 	Err  awserr.Error
 }
 
+type mockedGetParameterReturnPair struct {
+	Resp ssm.GetParametersByPathOutput
+	Err  awserr.Error
+}
+
 type mockedPutParameter struct {
 	ssmiface.SSMAPI
 	retVals   []mockedPutParameterReturnPair
 	callCount *int
 }
 
+type mockedGetParameter struct {
+	ssmiface.SSMAPI
+	retVal mockedGetParameterReturnPair
+}
+
 func (m mockedPutParameter) PutParameter(in *ssm.PutParameterInput) (*ssm.PutParameterOutput, error) {
 	// we have to increment a pointer to an integer because we can't update struct internal state here:
 	// ssmiface.SSMAPI requires PutParameter to be a receiver on a struct value, so this method gets a copy
 	// of the struct, not a pointer to a mockedPutParameter that we can mutate.
+	time.Sleep(time.Duration(1) * time.Microsecond) // We have to wait some amount to reliably validate timeouts
 	callCount := *m.callCount
 	resp := m.retVals[callCount]
 	*m.callCount = callCount + 1
 	return &resp.Resp, resp.Err
 }
 
-type TestCase struct {
+func (m mockedGetParameter) GetParametersByPath(input *ssm.GetParametersByPathInput) (*ssm.GetParametersByPathOutput, error) {
+	return &m.retVal.Resp, m.retVal.Err
+}
+
+type WriteSingleParamTestCase struct {
 	RetVals  []mockedPutParameterReturnPair
 	Expected error
 }
 
+type WriteToParameterStoreTestCase struct {
+	RetVals  []mockedPutParameterReturnPair
+	Timeout  time.Duration
+	Expected error
+}
+
+type ReadFromParameterStoreTestCase struct {
+	RetVal   mockedGetParameterReturnPair
+	Expected map[string]string
+}
+
+func TestReadFromParameterStore(t *testing.T) {
+	cases := []ReadFromParameterStoreTestCase{
+		{
+			RetVal: mockedGetParameterReturnPair{
+				Resp: ssm.GetParametersByPathOutput{
+					Parameters: []*ssm.Parameter{
+						&ssm.Parameter{
+							Name:  aws.String("name"),
+							Value: aws.String("value"),
+						},
+					},
+				},
+				Err: nil,
+			},
+			Expected: map[string]string{
+				"name": "value",
+			},
+		},
+	}
+
+	path := util.NewParameterStorePath("/path/parameter")
+
+	for i, c := range cases {
+		parameters := ReadFromParameterStore(*path, mockedGetParameter{retVal: c.RetVal})
+		if !reflect.DeepEqual(c.Expected, parameters) {
+			t.Fatalf("%v expected %v, got %v", i, c.Expected, parameters)
+		}
+	}
+}
+
 func TestWriteSingleParameter(t *testing.T) {
-	cases := []TestCase{
+	cases := []WriteSingleParamTestCase{
 		// happy case: no throttling, no error. smooth sailing!
 		{
 			RetVals: []mockedPutParameterReturnPair{
@@ -105,3 +164,54 @@ func TestWriteSingleParameter(t *testing.T) {
 		}
 	}
 }
+
+func TestWriteToParameterStore(t *testing.T) {
+	cases := []WriteToParameterStoreTestCase{
+		// Plenty of time to pretend to put parameters
+		{
+			RetVals: []mockedPutParameterReturnPair{
+				{
+					Resp: ssm.PutParameterOutput{
+						Tier:    aws.String("mock"),
+						Version: aws.Int64(1),
+					},
+					Err: nil,
+				},
+			},
+			Timeout:  time.Duration(1) * time.Minute,
+			Expected: nil,
+		},
+		// Not enough time to pretend to put parameters
+		{
+			RetVals: []mockedPutParameterReturnPair{
+				{
+					Resp: ssm.PutParameterOutput{
+						Tier:    aws.String("mock"),
+						Version: aws.Int64(1),
+					},
+					Err: nil,
+				},
+			},
+			Timeout:  time.Duration(0) * time.Minute,
+			Expected: errors.New("timeout"),
+		},
+	}
+
+	path := util.NewParameterStorePath("/path/")
+	for i, c := range cases {
+		callCount := 0
+		err := WriteToParameterStore(map[string]string{"path": "value"}, *path, c.Timeout, mockedPutParameter{retVals: c.RetVals, callCount: &callCount})
+		if c.Expected != nil {
+			if err == nil {
+				t.Fatalf("%d expected %d, got %d", i, c.Expected, err)
+			}
+			if err.Error() != c.Expected.Error() {
+				t.Fatalf("%d expected %d, got %d", i, c.Expected, err)
+			}
+		} else {
+			if err != nil {
+				t.Fatalf("%d expected %d, got %d", i, c.Expected, err)
+			}
+		}
+	}
+}