From cd1239de1c0ffeefbf1bc8f20db36d7aeeab6fd7 Mon Sep 17 00:00:00 2001
From: Allison Maheu <47192407+sendqueery@users.noreply.github.com>
Date: Mon, 24 Aug 2020 13:23:09 -0400
Subject: [PATCH] Remove unused function and add runner tests/mocks (#33)

---
 ssm/invocation/runner_test.go | 41 +++++++++++++++++++++++++++++++++++
 testing/ssm_mocks.go          | 37 ++++++++++++++++---------------
 2 files changed, 60 insertions(+), 18 deletions(-)

diff --git a/ssm/invocation/runner_test.go b/ssm/invocation/runner_test.go
index e811dd5..19c2755 100644
--- a/ssm/invocation/runner_test.go
+++ b/ssm/invocation/runner_test.go
@@ -1 +1,42 @@
 package invocation
+
+import (
+	"testing"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/service/ssm"
+	mocks "github.com/disneystreaming/ssm-helpers/testing"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGetResult(t *testing.T) {
+	assert := assert.New(t)
+	mockSvc := &mocks.MockSSMClient{}
+
+	successCmd, badCmd, mockInstance :=
+		aws.String("success-id"), aws.String("bad-id"), aws.String("i-123")
+
+	oc := make(chan *ssm.GetCommandInvocationOutput)
+	ec := make(chan error)
+
+	t.Run("valid ID", func(t *testing.T) {
+		go GetResult(mockSvc, successCmd, mockInstance, oc, ec)
+		select {
+		case result := <-oc:
+			assert.Equal("success-id", *result.CommandId)
+		case err := <-ec:
+			assert.Empty(err)
+		}
+
+	})
+
+	t.Run("invalid ID", func(t *testing.T) {
+		go GetResult(mockSvc, badCmd, mockInstance, oc, ec)
+		select {
+		case result := <-oc:
+			assert.Empty(result)
+		case err := <-ec:
+			assert.Error(err)
+		}
+	})
+}
diff --git a/testing/ssm_mocks.go b/testing/ssm_mocks.go
index 0ab45b0..b045ce6 100644
--- a/testing/ssm_mocks.go
+++ b/testing/ssm_mocks.go
@@ -2,7 +2,6 @@ package mocks
 
 import (
 	"fmt"
-	"os"
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/awserr"
@@ -56,32 +55,33 @@ func (m *MockSSMClient) TerminateSession(input *ssm.TerminateSessionInput) (outp
 }
 
 func (m *MockSSMClient) GetCommandInvocation(input *ssm.GetCommandInvocationInput) (output *ssm.GetCommandInvocationOutput, err error) {
-
-	if os.Getenv(fmt.Sprintf("%s-trycount", *input.InstanceId)) == "0" {
-		// Test waiting in cases where the invocation hasn't started yet
-		err = awserr.New("InvocationDoesNotExist", "The command ID and instance ID you specified did not match any invocations.\nVerify the command ID and the instance ID and try again.", err)
-		os.Setenv(fmt.Sprintf("%s-trycount", *input.InstanceId), "1")
-		return nil, err
-	} else if os.Getenv(fmt.Sprintf("%s-trycount", *input.InstanceId)) == "1" {
-		// Test waiting for the invocation to complete
-		output = &ssm.GetCommandInvocationOutput{
-			InstanceId:    input.InstanceId,
-			CommandId:     input.CommandId,
-			StatusDetails: aws.String("InProgress"),
+	var status string
+
+	switch *input.CommandId {
+	case "success-id":
+		status = "Success"
+	case "failed-id", "mixed-id":
+		status = "Failed"
+	case "pending-id":
+		status = "Pending"
+	case "bad-id":
+		return nil, awserr.New("InvalidCommandId", "InvalidCommandId", nil)
 		}
-		os.Setenv(fmt.Sprintf("%s-trycount", *input.InstanceId), "2")
-		return output, nil
-	} else {
+
 		output = &ssm.GetCommandInvocationOutput{
 			InstanceId:    input.InstanceId,
 			CommandId:     input.CommandId,
-			StatusDetails: aws.String("Success"),
+		StatusDetails: aws.String(status),
 		}
+
 		return output, nil
 	}
-}
 
 func (m *MockSSMClient) SendCommand(input *ssm.SendCommandInput) (output *ssm.SendCommandOutput, err error) {
+	if len(input.InstanceIds) > 0 && len(input.Targets) > 0 {
+		return nil, fmt.Errorf("Cannot specify instance IDs and SSM targets in same SendCommandInput")
+	}
+
 	// Mock our response from the SSM API
 	output = &ssm.SendCommandOutput{
 		Command: &ssm.Command{
@@ -89,6 +89,7 @@ func (m *MockSSMClient) SendCommand(input *ssm.SendCommandInput) (output *ssm.Se
 			DocumentName: input.DocumentName,
 			InstanceIds:  input.InstanceIds,
 			Parameters:   input.Parameters,
+			Targets:      input.Targets,
 		},
 	}
 	return output, nil