Stubbing AWS Service calls in Golang

I recently switched to Golang for my language of choice. (In my previous blog you can read why.) But I am also a big fan of test driven development. With Python you have a stubber that helps you mock the AWS API. So how do you do this in Golang? In this blog I will share my experience so far.

Use dependency injection

My first experiment was with dependency injection. I used the following code to do this:

package main

import (
    "context"
    "github.com/aws/aws-sdk-go-v2/config"
    "github.com/aws/aws-sdk-go-v2/service/s3"
    "time"
    "log"
    "os"
)

type Request struct {}
type Response struct {}

type Lambda struct {
    s3Client *s3.Client
}

func New() (*Lambda, error) {
    cfg, err := config.LoadDefaultConfig(context.TODO())

    m := new(Lambda)
    m.SetS3Client(s3.NewFromConfig(cfg))
    return m, err
}

func (x *Lambda) SetS3Client(client *s3.Client) {
    x.s3Client = client
}

func (x *Lambda) Handler(ctx context.Context, request Request) (Response, error) {
    // Your lambda code goes here
}

In your tests you could now use it as followed:

package main

import (
    "context"
    "github.com/aws/aws-sdk-go-v2/service/s3"
    "testing"
    "time"
    "log"
    "os"
)

type mockS3Client struct {
    s3.Client
    Error error
}

func (m *mockS3Client) PutObject(input *s3.PutObjectInput) (*s3.PutObjectOutput, error) {
    return &s3.PutObjectOutput{}, nil
}

func TestHandler(t *testing.T) {
    lambda := New()
    lambda.SetS3Client(&mockS3Client{})
    var ctx = context.Background()
    var event Request

    t.Run("Invoke Handler", func(t *testing.T) {
        response, err := lambda.Handler(ctx, event)

        // Perform Assertions
    })
}

We inject a mocked object that acts as the client used to perform the API calls. With this approach I could now write some tests. But I realized that this approach creates another problem. For example, what if you have 2 API calls that perform a PutObject call. In this example I return an empty PutObjectOutput. But I want to test more than one scenarios, so how do you control this behavior in your mocked object?

Using a stubber

So I did some more research and I found the awsdocs/aws-doc-sdk-examples repo. This repository used a testtools module. So I started an experiment to see how I could use this module. I refactored the code as followed:

package main

import (
    "context"
    "github.com/aws/aws-sdk-go-v2/aws"
    "github.com/aws/aws-sdk-go-v2/service/s3"
)

type Request struct {}
type Response struct {}

type Lambda struct {
    ctx      context.Context
    s3Client *s3.Client
}

func New(cfg aws.Config) *Lambda {
    m := new(Lambda)
    m.s3Client = s3.NewFromConfig(cfg)
    return m
}

func (x *Lambda) Handler(ctx context.Context, request Request) (Response, error) {
    // Your lambda code goes here
    return Response{}, nil
}

I added a cfg parameter to the New method, so I also need to pass this in my main method.

package main

import (
    "context"
    "github.com/aws/aws-lambda-go/lambda"
    "github.com/aws/aws-sdk-go-v2/config"
    "log"
)

func main() {
    cfg, err := config.LoadDefaultConfig(context.TODO())
    if err != nil {
        log.Printf("error: %v", err)
        return
    }
    lambda.Start(New(cfg).Handler)
}

The test itself looks like this:

package main

import (
    "bytes"
    "context"
    "encoding/json"
    "errors"
    "github.com/aws/aws-sdk-go-v2/aws"
    "github.com/aws/aws-sdk-go-v2/service/s3"
    "github.com/awsdocs/aws-doc-sdk-examples/gov2/testtools"
    "io"
    "os"
    "strings"
    "testing"
)

func TestHandler(t *testing.T) {
    var ctx = context.Background()
    var event Request

    t.Run("Upload a file to S3", func(t *testing.T) {
        stubber := testtools.NewStubber()
        lambda := New(*stubber.SdkConfig)

        stubber.Add(testtools.Stub{
            OperationName: "PutObject",
            Input: &s3.PutObjectInput{
                Bucket: aws.String("my-sample-bucket"),
                Key:    aws.String("my/object.json"),
                Body:   bytes.NewReader([]byte{}),
            },
            Output: &s3.PutObjectOutput{},
        })

        response, err := lambda.Handler(ctx, event)
        testtools.ExitTest(stubber, t)

        // Perform Assertions
    })
}

As you can see, we now moved the mock in the test itself. This enables you to let the AWS API react based on your test. The biggest advantage is that it’s encapsulated in the test itself. For example, If you want to add a scenario where the PutObject call failed you add the following:

t.Run("Fail on upload", func(t *testing.T) {
    stubber := testtools.NewStubber()
    lambda := New(*stubber.SdkConfig)
    raiseErr := &testtools.StubError{Err: errors.New("ClientError")}

    stubber.Add(testtools.Stub{
        OperationName: "PutObject",
        Input: &s3.PutObjectInput{
            Bucket: aws.String("my-sample-bucket"),
            Key:    aws.String("my/object.json"),
            Body:   bytes.NewReader([]byte{}),
        },
        Error: raiseErr,
    })

    _, err := lambda.Handler(ctx, event)
    testtools.VerifyError(err, raiseErr, t)
    testtools.ExitTest(stubber, t)
})

The testtools.VerifyError(err, raiseErr, t) definition will confirm if the error is indeed passed along. The testtools.ExitTest(stubber, t) definition will fail the test if a stub that you added was not called. You can use this to confirm if all expected API calls where indeed executed.

In some cases you want to ignore certain fields in your Input. You can add a list of IgnoreFields: []string{"MyField"} to your stubber. This is useful if you do not have direct control over what is send.

Conclusion

The testtool is a good replacement of the stubber I used in Python. It allows you to encapsulate scenario data in your test. Avoiding hard to maintain mock objects. The testtool works from the configuration, so you don’t need to stub every client. Resulting in less code that is needs to test your implementation.

Photo by Klaus Nielsen