How to Unit Test a Lambda Implemented with AWS SDK for Go V2

Question

Given the following simple lambda written in Go that just returns a table description...

package main

import (
    "context"
    "encoding/json"
    "fmt"
    "log"
    "net/http"
    "os"
    "strings"

    "github.com/aws/aws-lambda-go/events"
    "github.com/aws/aws-lambda-go/lambda"
    "github.com/aws/aws-sdk-go-v2/aws"
    "github.com/aws/aws-sdk-go-v2/config"
    "github.com/aws/aws-sdk-go-v2/service/dynamodb"
    "go.uber.org/zap"
)

var (
    dynamoDBTableName = aws.String(os.Getenv(EnvDynamoDBTableName))

    logger = func() *zap.Logger {
        l, err := zap.NewProduction()
        if err != nil {
            log.Printf("failed to create zap logger: %v", err)
        }
        return l
    }()
)

func handler(ctx context.Context, req events.APIGatewayProxyRequest) 
    (events.APIGatewayProxyResponse, error) {

    defer logger.Sync()

    resp := events.APIGatewayProxyResponse{}

    cfg, err := config.LoadDefaultConfig(ctx)
    if err != nil {
        logger.Error("failed to load AWS config", zap.Error(err))
        return resp, fmt.Errorf("failed to load AWS config: %w", err)
    }

    svc := dynamodb.NewFromConfig(cfg)

    // fake logic
    t, err := svc.DescribeTable(ctx, &dynamodb.DescribeTableInput{TableName: dynamoDBTableName})
    if err != nil {
        logger.Error("failed to describe table", zap.String("table-name", *dynamoDBTableName), zap.Error(err))
    }
    var sb strings.Builder
    enc := json.NewEncoder(&sb)
    err = enc.Encode(t.Table)
    if err != nil {
        logger.Error("failed to JSON encode response", zap.Error(err))
    }
    resp.Body = sb.String()
    resp.StatusCode = http.StatusOK

    return resp, nil
}

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

... how do I unit test it locally? With the old SDK it was possible to use dependency injection like this:

type deps struct 
    svc dynamodbiface.DynamoDBAPI
    table string
}

func (d *deps) handler(ctx context.Context, req events.APIGatewayProxyRequest) 
    (events.APIGatewayProxyResponse, error) {
    
    ...
}

func main() {
   s := session.Must(session.NewSession())
   d := deps {
       svc: dynamodb.New(s),
       table: dynamoDBTableName,
   }

   lambda.Start(d.handler)
}

How do I test a lambda written with the new AWS SDK for Go V2 given that I need the context to load the config required by dynamodb.NewFromConfig?

user_17013248 · Answer

First make your handler a humble object, so we can "skip" testing it:

func handler(ctx context.Context, req events.APIGatewayProxyRequest) 
    (events.APIGatewayProxyResponse, error) {
    dynamoWrapper := &RealDynamoWrapper{}
    proxyController := &ProxyController{DynamoWrapper: dynamoWrapper}
    return proxyController.proxy(ctx, req)

The idea is to make the handler function humble by making it delegate all the complexity to the proxyController.

Now, lets consider the ProxyController to be under test, we need to define it first:

type ProxyController struct {
  dynamoWrapper DynamoWrapper
}

func(controller *ProxyController) Proxy(ctx context.Context, req events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
  // this is where your implementation lives that you need to mock stuff for

  // We are mocking this line from your question
  svc := controller.dynamoWrapper.NewFromConfig(...)

  // do other stuff and then return something
  return events.APIGatewayProxyResponse{}, nil
}

You can see I'm going to depend on a wrapped version of dynamo, which looks like this:

type DynamoWrapper interface {
  NewFromConfig(cfg aws.Config, optFns ...func(*Options)) *Client
}

Now the real implementation of this wrapper, the one referenced above as RealDynamoWrapper will make the call as you do to the sdk. For our test though, we want a mock implementation:

type mockDynamoWrapper struct {
  NewFromConfigFunc func(aws.Config, ...func(*Options)) *Client
}

func(dynamoWrapper *mockDynamoWrapper) NewFromConfig(cfg aws.Config, optFns ...func(*Options)) *Client {
  return dynamoWrapper.NewFromConfigFunc(cfg, optFns...)
}

Lastly in your test you can now mock the dynamo call:

func TestProxyController(t *testing.T) {
  // given
  dynamoWrapper := &mockDynamoWrapper{}
  proxyController := &ProxyController{DynamoWrapper: mockDynamoWrapper}
  request := events.APIGatewayProxyRequest{}
  dynamoWrapper.NewFromConfigFunc = func(aws.Config, ...func(*Options)) *Client {
    // setup your mock function to do whatever you want
  }

  // when
  proxyController.proxy(context.Background, request)

  // then do your asserts
}

How to Unit Test a Lambda Implemented with AWS SDK for Go V2

1 Answers