From 8774e95325f46c909ee19cb0395ced01450f2658 Mon Sep 17 00:00:00 2001
From: Karri Balk <kbalk@users.noreply.github.com>
Date: Mon, 29 Mar 2021 12:12:24 -0400
Subject: [PATCH] Add a InvokceFunctionWithParamsE() func

---
 modules/aws/lambda.go                     | 26 +++++++++++++++--------
 test/terraform_aws_lambda_example_test.go | 19 +++++++++++++++--
 2 files changed, 34 insertions(+), 11 deletions(-)

diff --git a/modules/aws/lambda.go b/modules/aws/lambda.go
index ac4fe8bc4..bee886944 100644
--- a/modules/aws/lambda.go
+++ b/modules/aws/lambda.go
@@ -14,12 +14,11 @@ import (
 type LambdaOptions struct {
 	// 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.
+	//    * InvocationTypeRequestResponse (default) - Invoke the function
+	//      synchronously.  Keep the connection open until the function
+	//      returns a response or times out.
+	//    * InvocationTypeDryRun - Validate parameter values and verify
+	//      that the user or role has permission to invoke the function.
 	InvocationType *string `enum:"lambda.InvocationType"`
 
 	// Lambda function input; will be converted to JSON.
@@ -53,7 +52,7 @@ func InvokeFunction(t testing.TestingT, region, functionName string, payload int
 // InvokeFunctionE invokes a lambda function.
 func InvokeFunctionE(t testing.TestingT, region, functionName string, payload interface{}) ([]byte, error) {
 	input := &LambdaOptions{Payload: &payload}
-	out, err := InvokeFunctionWithParams(t, region, functionName, input)
+	out, err := InvokeFunctionWithParamsE(t, region, functionName, input)
 	if err != nil {
 		return nil, err
 	}
@@ -67,8 +66,17 @@ func InvokeFunctionE(t testing.TestingT, region, functionName string, payload in
 
 // InvokeFunctionWithParams invokes a lambda function using parameters
 // supplied in the LambdaOptions struct and returns values in a LambdaOutput
-// struct.
-func InvokeFunctionWithParams(t testing.TestingT, region, functionName string, input *LambdaOptions) (*LambdaOutput, error) {
+// struct.  Checks for failure using "require".
+func InvokeFunctionWithParams(t testing.TestingT, region, functionName string, input *LambdaOptions) *LambdaOutput {
+	out, err := InvokeFunctionWithParamsE(t, region, functionName, input)
+	require.NoError(t, err)
+	return out
+}
+
+// InvokeFunctionWithParamsE invokes a lambda function using parameters
+// supplied in the LambdaOptions struct and returns values in a LambdaOutput
+// struct and the error.
+func InvokeFunctionWithParamsE(t testing.TestingT, region, functionName string, input *LambdaOptions) (*LambdaOutput, error) {
 	lambdaClient, err := NewLambdaClientE(t, region)
 	if err != nil {
 		return nil, err
diff --git a/test/terraform_aws_lambda_example_test.go b/test/terraform_aws_lambda_example_test.go
index 9f4b9acc2..e486dda1d 100644
--- a/test/terraform_aws_lambda_example_test.go
+++ b/test/terraform_aws_lambda_example_test.go
@@ -104,13 +104,28 @@ func TestTerraformAwsLambdaWithParamsExample(t *testing.T) {
 	// test function will not be checking the payload.
 	invocationType := lambda.InvocationTypeDryRun
 	input := &aws.LambdaOptions{InvocationType: &invocationType}
-	out, err := aws.InvokeFunctionWithParams(t, awsRegion, functionName, input)
+	out := aws.InvokeFunctionWithParams(t, awsRegion, functionName, input)
 
 	// With "DryRun", there's no message in the output, but there is
 	// a status code which will have a value of 204 for a successful
 	// invocation.
-	require.Nil(t, err)
 	assert.Equal(t, int(*out.StatusCode), 204)
+
+	// Invoke the function, this time causing the Lambda to error and
+	// capturing the error.
+	invocationType = lambda.InvocationTypeRequestResponse
+	input = &aws.LambdaOptions{
+		InvocationType: &invocationType,
+		Payload:        ExampleFunctionPayload{ShouldFail: true, Echo: "hi!"},
+	}
+	out, err := aws.InvokeFunctionWithParamsE(t, awsRegion, functionName, input)
+
+	// No error in the invocation as Lambda was found and executed.
+	require.Nil(t, err)
+	assert.Equal(t, int(*out.StatusCode), 200)
+
+	// Make sure the function-specific error comes back
+	assert.Contains(t, string(out.Payload), "Failed to handle")
 }
 
 type ExampleFunctionPayload struct {