From 6d632a920bf71534d0967f20b6cfbd6630577abd Mon Sep 17 00:00:00 2001
From: Karri Balk <kbalk@users.noreply.github.com>
Date: Mon, 29 Mar 2021 11:34:47 -0400
Subject: [PATCH] Use enum for InvocationType

---
 modules/aws/lambda.go                     | 26 ++++-------------------
 test/terraform_aws_lambda_example_test.go | 16 ++------------
 2 files changed, 6 insertions(+), 36 deletions(-)

diff --git a/modules/aws/lambda.go b/modules/aws/lambda.go
index 370b5b8fc..ac4fe8bc4 100644
--- a/modules/aws/lambda.go
+++ b/modules/aws/lambda.go
@@ -2,7 +2,6 @@ package aws
 
 import (
 	"encoding/json"
-	"errors"
 	"fmt"
 
 	"github.com/aws/aws-sdk-go/service/lambda"
@@ -13,14 +12,15 @@ import (
 // LambdaOptions contains additional parameters for InvokeFunctionWithParams().
 // It contains a subset of the fields found in the lambda.InvokeInput struct.
 type LambdaOptions struct {
-	// InvocationType can be one of "RequestResponse" or "DryRun".
+	// InvocationType can be one of lambda.InvocationTypeRequestResponse
+	// or lambda.InvocationTypeDryRun.
 	//    * RequestResponse (default) - Invoke the function synchronously.
 	//    Keep the connection open until the function returns a response
 	//    or times out.
 	//
 	//    * DryRun - Validate parameter values and verify that the user or
 	//    role has permission to invoke the function.
-	InvocationType *string
+	InvocationType *string `enum:"lambda.InvocationType"`
 
 	// Lambda function input; will be converted to JSON.
 	Payload interface{}
@@ -74,27 +74,9 @@ func InvokeFunctionWithParams(t testing.TestingT, region, functionName string, i
 		return nil, err
 	}
 
-	// Verify the InvocationType is one of the allowed values and report
-	// an error if its not.  By default the InvocationType will be
-	// "RequestResponse".
-	invocationType := lambda.InvocationTypeRequestResponse
-	if input.InvocationType != nil {
-		switch *input.InvocationType {
-		case
-			lambda.InvocationTypeRequestResponse,
-			lambda.InvocationTypeDryRun:
-			invocationType = *input.InvocationType
-		default:
-			msg := fmt.Sprintf("LambdaOptions.InvocationType, if specified, must either be \"%s\" or \"%s\"",
-				lambda.InvocationTypeRequestResponse,
-				lambda.InvocationTypeDryRun)
-			return &LambdaOutput{FunctionError: &msg}, errors.New(msg)
-		}
-	}
-
 	invokeInput := &lambda.InvokeInput{
 		FunctionName:   &functionName,
-		InvocationType: &invocationType,
+		InvocationType: input.InvocationType,
 	}
 
 	if input.Payload != nil {
diff --git a/test/terraform_aws_lambda_example_test.go b/test/terraform_aws_lambda_example_test.go
index ca2b78493..9f4b9acc2 100644
--- a/test/terraform_aws_lambda_example_test.go
+++ b/test/terraform_aws_lambda_example_test.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"testing"
 
+	"github.com/aws/aws-sdk-go/service/lambda"
 	"github.com/gruntwork-io/terratest/modules/aws"
 	"github.com/gruntwork-io/terratest/modules/random"
 	"github.com/gruntwork-io/terratest/modules/terraform"
@@ -101,7 +102,7 @@ func TestTerraformAwsLambdaWithParamsExample(t *testing.T) {
 	// Call InvokeFunctionWithParms with an InvocationType of "DryRun".
 	// A "DryRun" invocation does not execute the function, so the example
 	// test function will not be checking the payload.
-	invocationType := "DryRun"
+	invocationType := lambda.InvocationTypeDryRun
 	input := &aws.LambdaOptions{InvocationType: &invocationType}
 	out, err := aws.InvokeFunctionWithParams(t, awsRegion, functionName, input)
 
@@ -110,19 +111,6 @@ func TestTerraformAwsLambdaWithParamsExample(t *testing.T) {
 	// invocation.
 	require.Nil(t, err)
 	assert.Equal(t, int(*out.StatusCode), 204)
-
-	// Call InvokeFunctionWithParams with a LambdaOptions struct with an
-	// unsupported InvocationType.  The function should fail.
-	invocationType = "Event"
-	input = &aws.LambdaOptions{
-		InvocationType: &invocationType,
-		Payload:        ExampleFunctionPayload{ShouldFail: false, Echo: "hi!"},
-	}
-	out, err = aws.InvokeFunctionWithParams(t, awsRegion, functionName, input)
-	require.NotNil(t, err)
-	msg := "LambdaOptions.InvocationType, if specified, must either be \"RequestResponse\" or \"DryRun\""
-	assert.Contains(t, err.Error(), msg)
-	assert.Contains(t, *out.FunctionError, msg)
 }
 
 type ExampleFunctionPayload struct {