Hands On: Writing and testing your Go Lambdas

In this section we’ll add a new function to our application. This new function’s purpose will be to capture the shoutout that a user submits through the Slack slash command. Recall that in our architecture diagram, one of the functions in our state machine was writing to a DynamoDB table. That’s the function you’re going to write, complete with a unit test.

The steps involved are:

1. Writing the test
2. Writing the code to pass the test
3. Update our SAM template to define the new function
4. Update our build process to include the new function
5. Deploying and debugging as needed

Writing tests (Hands On)

Setup

To get the most out of this hands-on exercise, you will switch to the hands-on-saver branch of the project's repository:

$ git fetch --all
$ git checkout hands-on-saver

The version of the project in this branch omits the saving (to DynamoDB) behavior that our solution relies on. You are going to replicate that here.

Let's Go!

We’re going to write a unit test to verify the behavior of our code in isolation from the test of the application. Let’s start by creating a saver_test.go file and adding the following code to it:

package shoutouts_test

import (
	"context"
	"errors"
	"testing"

	"github.com/aws/aws-sdk-go/aws/request"

	"github.com/stretchr/testify/assert"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/dynamodb"
	"github.com/jboursiquot/shoutouts"
)

func TestSaver(t *testing.T) {
	cases := []struct {
		scenario string
		shoutout *shoutouts.Shoutout
		ddb      *mockDynamoDB
	}{
		{
			scenario: "api error",
			shoutout: shoutouts.New(),
			ddb: &mockDynamoDB{
				putOut: nil,
				err:    errors.New("api error"),
			},
		},
		{
			scenario: "successful put",
			shoutout: shoutouts.New(),
			ddb: &mockDynamoDB{
				putOut: &dynamodb.PutItemOutput{},
				err:    nil,
			},
		},
		{
			scenario: "successful query",
			shoutout: shoutouts.New(),
			ddb: &mockDynamoDB{
				queryOut: &dynamodb.QueryOutput{},
				err:      nil,
			},
		},
	}

	for _, c := range cases {
		t.Run(c.scenario, func(t *testing.T) {
			s := shoutouts.NewSaver(c.ddb)

			if c.ddb.err != nil {
				assert.Error(t, s.Save(context.Background(), c.shoutout))
			}

			if c.ddb.err == nil {
				assert.NoError(t, s.Save(context.Background(), c.shoutout))
			}
		})
	}
}

type mockDynamoDB struct {
	putOut   *dynamodb.PutItemOutput
	queryOut *dynamodb.QueryOutput
	err      error
}

func (m *mockDynamoDB) PutItemWithContext(ctx aws.Context, item *dynamodb.PutItemInput, opts ...request.Option) (*dynamodb.PutItemOutput, error) {
	return m.putOut, m.err
}

func (m *mockDynamoDB) QueryWithContext(aws.Context, *dynamodb.QueryInput, ...request.Option) (*dynamodb.QueryOutput, error) {
	return m.queryOut, m.err
}

There are a few things going on here so let’s break it down. First, notice the package declaration of package shoutouts_test which puts this code outside of package shoutouts where our actual behavior resides. We do this deliberately sometimes as a means of testing from the outside of the package under test which eliminates the temptation to use non-exported values.

Next, let’s look at the TestSaver function. It uses a technique called table-driven testing which lets you specify multiple test cases in a single variable that we then iterate on during our tests. In this test, we’ve set up three scenarios, to keep things simple. The first simulates an error communicating with our data storage API while the second simulates a successful attempt to do so. The third is used for testing queries but that's for a later discussion.

We iterate over the cases and make use of the subtest feature of the testing package to run each scenario. For each case, we initialize a shoutouts.Saver, through the NewSaver of the same package. That type and function do not yet exist but that we’ll create them soon. The next two sets of conditions simply check whether we’re expecting an error from calling the Save method on the saver or not.

If you run the test at this stage, you should get something like the following error:

/Users/jboursiquot/dev/go/src/github.com/jboursiquot/shoutouts/saver_test.go:41:9: undefined: shoutouts.NewSaver
FAIL	github.com/jboursiquot/shoutouts [build failed]
Error: Tests failed.

This makes sense because the shoutouts.NewSaver function is not yet defined. Let’s do so now. Add a saver.go file to the root of the project with the following content:

package shoutouts

import (
	"context"
	"fmt"
	"os"

	"github.com/aws/aws-sdk-go/aws/request"

	"github.com/aws/aws-sdk-go/aws"

	"github.com/aws/aws-sdk-go/service/dynamodb"
	"github.com/aws/aws-sdk-go/service/dynamodb/dynamodbattribute"
)

// DynamoDBPuter is the minimal interface needed to store a shoutout.
type DynamoDBPuter interface {
	PutItemWithContext(aws.Context, *dynamodb.PutItemInput, ...request.Option) (*dynamodb.PutItemOutput, error)
}

// NewSaver returns a new saver.
func NewSaver(c DynamoDBPuter) *Saver {
	return &Saver{ddb: c}
}

// Saver is a shoutout saver.
type Saver struct {
	ddb DynamoDBPuter
}

// Save saves a Shoutout
func (s *Saver) Save(ctx context.Context, shoutout *Shoutout) error {
	item, err := dynamodbattribute.MarshalMap(shoutout)
	if err != nil {
		return fmt.Errorf("failed to marshal shoutout for storage: %s", err)
	}

	input := &dynamodb.PutItemInput{
		Item:      item,
		TableName: aws.String(os.Getenv("TABLE_NAME")),
	}

	if _, err = s.ddb.PutItemWithContext(ctx, input); err != nil {
		return fmt.Errorf("failed to save shoutout: %s", err)
	}

	return nil
}

As with the test file we looked at earlier, note the package declaration. This one is shoutouts and not shoutouts_test as we discussed earlier.

Next, we see a DynamoDBPuter interface defined with a single method.

// DynamoDBPuter is the minimal interface needed to store a shoutout.
type DynamoDBPuter interface {
	PutItemWithContext(aws.Context, *dynamodb.PutItemInput, ...request.Option) (*dynamodb.PutItemOutput, error)
}

The reason for this interface is twofold: first, we know we need to have a client that can talk to DynamoDB eventually and so will need to pass that in at some point when our function gets executed and second, we need a way of unit testing our code without talking to a live DynamoDB API. By making use of an interface, both our “real” code and our test code can pass in an implementation for this interface. The reason for the single method is because that’s the only one that we need for this code to work. There’s no point in mocking the entire DynamoDB API when we’re only using one function in it.

Our "mock" is more of a "stub" here but we’re using the term "mock" to mean that we’re not actually talking to DynamoDB.

Our Saver struct needs a DynamoDB client to function and that's there the DynamoDBPuter interface comes in. When we initialize a new Saver using NewSaver(c DynamoDBPuter), we'll pass in something that satisfies the DynamoDBPuter interface whether that argument be a real DynamoDB client or a mock implementation (see mockDynamoDB in saver_test.go).

The Saver struct's Save method is where the action happens. In this method, we marshal the incoming shoutout into a map of DynamoDB attribute values, create an input struct with this map and the table name we want to write to, and finally use the DynamoDB API client we initialized our Saver type with to make the call to save the item.

If we run our test again now, we should see them pass:

ok  	github.com/jboursiquot/shoutouts
Success: Tests passed.

Next up, we need to write the Lambda function code that makes use of our newly coded save behavior.

Create cmd/saver/main.go and in that file put the following code:

package main

import (
	"context"

	"github.com/aws/aws-lambda-go/lambda"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/dynamodb"
	"github.com/aws/aws-xray-sdk-go/xray"
	"github.com/jboursiquot/shoutouts"
)

var sess *session.Session
var ddb *dynamodb.DynamoDB

func init() {
	sess = session.Must(session.NewSession())
	ddb = dynamodb.New(sess)
}

func handler(ctx context.Context, shoutout *shoutouts.Shoutout) (*shoutouts.Shoutout, error) {
	err := shoutouts.NewSaver(ddb).Save(ctx, shoutout)
	if err != nil {
		xray.AddError(ctx, err)
	}
	return shoutout, err
}

func main() {
	lambda.Start(handler)
}

It’s a relatively short bit of code. Our main function should look familiar. It simply uses the AWS Lambda SDK’s Start function to pass in a handler which we define as receiving a context.Context and a *shoutout.Shoutout value.

Remember that this Lambda function will be used in a state machine which will be passing in as input the shoutout that was originally captured and placed in the queue for processing. That JSON payload will automatically be un-marshaled into a Shoutout value for us when our saver function is invoked.

Inside the handler we’re returning that shoutout that was passed in and the result of the Saver type's Save call. That actual DynamoDB client needed for the Saver to work for real is initialized as part of the init function also in this code.

If you’re wondering why we set the variables sess and ddb as package globals, it’s because Lambda will "freeze" these variables for a period of time so that subsequent invokations of the Function initialize faster because the global state will be unfrozen and reused (see AWS Lambda Execution Context for more on this).

Another key thing to note here is that this time, we’re initializing a saver with shoutouts.NewSaver with an actual DynamoDB client and not a mock like we were doing in our tests. Since the real client will implicitly satisfy the DynamoDBPuter interface containing the PutItemWithContext method, everything will just snap into place.

That’s all there is to this Lambda function which makes use of our shoutout-saving behavior. Now that we have our function and feeling confident of its correctness based on our tests, in the next section, we’ll see how to include it as part of our orchestration.