Stubbing AWS Service calls in Golang

#go #stubbing #aws #mock

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? I found 2 ways to do this. One via dependency injection and one via a stubber. In this blog I will share my experience so far.

Using 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) {
    data := []byte("Hello World")

    _, err := x.s3Client.PutObject(x.ctx, &s3.PutObjectInput{
        Bucket: aws.String("my-bucket"),
        Key:    aws.String("my-object-key"),
        Body:   bytes.NewReader(data),
    })

    return Response{}, err
}

As you can see, I use the SetS3Client method as a setter. It allows me to set the client from the outside. This is useful when doing unit tests. You can use this in your tests 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) {
    data := []byte("Hello World")

    _, err := x.s3Client.PutObject(x.ctx, &s3.PutObjectInput{
        Bucket: aws.String("my-bucket"),
        Key:    aws.String("my-object-key"),
        Body:   bytes.NewReader(data),
    })

    return Response{}, err
}

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 main advantage of the stubber is that you validate the Input out of the box. But in some cases you want to ignore certain fields. For example, if you use a timestamp in your Key. Or the actual Body of the object you are uploading.

You can ignore these fields by setting the IgnoreFields. (Example: IgnoreFields: []string{"Key", "Body"})

The second thing I like about using the stubber is the VerifyError method. This will verify if the error is returned that you raised in your test scenario.

The last one is the ExitTest method. This will make sure that all the defined stubs are actually called. In other words, if you still have a stub. Your test will fail because there is still a uncalled stub.

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