Stubbing AWS Service calls in Golang
Tags: go, stubbing, aws, mock
Published: 2023-08-15 07:51:15

I recently switched to Golang for my language of choice. (In my previous [blog](https://xebia.com/blog/using-golang-for-your-aws-lambda-functions/) you can read why.) But I am also a big fan of test driven development. With Python you have a [stubber](https://botocore.amazonaws.com/v1/documentation/api/latest/reference/stubber.html) 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:

```go
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:

```go
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](https://github.com/awsdocs/aws-doc-sdk-examples/tree/main) 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:

```go
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.

```go
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:

```go
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:

```go
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](https://www.pexels.com/photo/brave-doctor-in-flying-superhero-cape-with-fist-stretched-6303759/)