Testing with AWS SDK for Go V2 without interface mocks

I found a good way to write unit tests in AWS SDK for Go V2 without having to mock the interface, so I put together some patterns you can use.

---

Overview

When writing unit tests for an application using AWS SDK for Go V2, you probably create an AWS SDK Client interface and mock the test code by creating stubs, using gomock, etc.

However, I think it's hard to create an SDK Client interface for every single test, so I'll write a test method that doesn't require mocking with an interface.

---

Middleware

Contents

The method is "Mock the SDK response using the middleware provided by AWS SDK for Go V2".

In other words, you can use the Client of the SDK as is for testing.

See doc.

However, in the next chapter of the official documentation("Testing"), it is recommended that testing be done by mocking the interface.

How it works

From request to response in AWS SDK for Go V2, processing is based on each step of the stack as shown in the figure here.

Stack Step	Description
Initialize	Prepares the input, and sets any default parameters as needed.
Serialize	Serializes the input to a protocol format suitable for the target transport layer.
Build	Attach additional metadata to the serialized input, such as HTTP Content-Length.
Finalize	Final message preparation, including retries and authentication (SigV4 signing).
Deserialize	Deserialize responses from the protocol format into a structured type or error.

You can mock this by customizing the input/output of each stack step.

---

Usage (outline)

middleware function definition

As a rough description of how to use the middleware, define the behavior of the middleware using the middleware functions provided for each stack, and pass them to the SDK to apply the middleware.

The following is an example of a middleware function called InitializeMiddlewareFunc that customizes the Client request parameters (xxInput) in the first Initialize step of the above steps.

GetObjectInput call, if there is no Bucket in the request parameters, specify "my-default-bucket" as Bucket.

var defaultBucket = middleware.InitializeMiddlewareFunc("DefaultBucket", func(
    ctx context.Context, in middleware.InitializeInput, next middleware.InitializeHandler,
) (
    out middleware.InitializeOutput, metadata middleware.Metadata, err error,
) {
    // Type switch to check if the input is s3.GetObjectInput, if so and the bucket is not set, populate it with
    // our default.
    switch v := in.Parameters.(type) {
    case *s3.GetObjectInput:
        if v.Bucket == nil {
            v.Bucket = aws.String("my-default-bucket")
        }
    }

    // Middleware must call the next middleware to be executed in order to continue execution of the stack.
    // If an error occurs, you can return to prevent further execution.
    return next.HandleInitialize(ctx, in)
})

How to pass middleware functions

There are two specific ways to pass middleware functions to the SDK.

• (1) Inject as an Option at the time of config definition (SDK)
• (2) Inject as an Option at the time of a specific operation (API call)

The difference between the two is "how and when to inject options," but specifically they have the following characteristics.

(1) Inject as an Option at the time of (SDK) config definition

This one is passed to APIOptions in the config used when generating each Client.

Also, when passing it, you can attach the middleware to each step by wrapping the function you just defined as a function that returns stack.Initialize.Add(defaultBucket, middleware.Before) as follows.

cfg, err := config.LoadDefaultConfig(context.TODO())
if err != nil {
    // handle error
}

cfg.APIOptions = append(cfg.APIOptions, func(stack *middleware.Stack) error {
    // Attach the custom middleware to the beginning of the Initialize step
    return stack.Initialize.Add(defaultBucket, middleware.Before)
})

client := s3.NewFromConfig(cfg)

Features include the following

• Common definitions can be made for each client and each operation.
  • As described later in (2), the method (1) here can also be used to branch processing for each operation.
• Because there is no need to modify the arguments of each API function, it is easy to use in testing.
  • In unit tests, generate a Client using a mock config with Option passed, and use it in tests.
  • See (2) below.

(2) Inject as an Option during a specific operation (API call)

This is not a config as in the previous example, but is passed as a function option when an API call is actually made from the Client.

When passing it, define a function to set options.APIOptions to options.APIOptions (complicated...) as the third argument of the API function, which is the same as the previous one, wrap function to attach middleware to the stack.

// registerDefaultBucketMiddleware registers the defaultBucket middleware with the provided stack.
func registerDefaultBucketMiddleware(stack *middleware.Stack) error {
    // Attach the custom middleware to the beginning of the Initialize step
    return stack.Initialize.Add(defaultBucket, middleware.Before)
}

// ...

cfg, err := config.LoadDefaultConfig(context.TODO())
if err != nil {
    // handle error
}

client := s3.NewFromConfig(cfg)

object, err := client.GetObject(context.TODO(), &s3.GetObjectInput{
    Key: aws.String("my-key"),
}, func(options *s3.Options) {
    // Register the defaultBucketMiddleware for this operation only
    options.APIOptions = append(options.APIOptions, registerDefaultBucketMiddleware)
})

Characteristics include the following

• Easy control of middleware for each API call (type of operation)
• You may need to rewrite how you call the SDK in production code to be middleware-aware, since you will be passing it directly to the third argument of the actual API call function (client.GetObject above)
  • If you use DI for the options function from the caller, there's no problem.

Regarding the latter point, it is fine if you are in an environment where you usually "DI" the options function as the third argument of API functions, but if you write the options function directly, or if you are in an environment where "I never fiddle with options...", you may omit the third argument and pass only two arguments. I guess there are some people who do this. (Yes, that's me)

In that case, the actual code needs to be rewritten to "set the options function to the third argument of each functions such as client.GetObject" and "rewrite to the DI so that the middleware can be successfully passed only in the test environment, without actually passing a function with content".

As for the former (Easy control of middleware for each API call (type of operation)), even when an option is injected in config as in (1), the API name (operation name) can be obtained in the middleware function, so the behavior can be changed for each API by writing branch processing based on it. (This will be introduced in the second pattern in the next chapter, "Mock Patterns".)

---

Mock Pattern

Now, at length, let's write the test code to actually mock using Middleware.

Here, we inject the middleware with config option, which is (1) above.

(2) is more flexible, but as mentioned earlier (in my environment) "I had to rewrite the code for testing ", so I chose (1).

I also wrote it in three separate pattern examples.

(1) Pattern with only one type of SDK API call in one test case

Overall view

This is the simplest case, where only one type of SDK API call is made in one test case.

Let's say you have code that uses an S3 Client like this, and you want to test the DeleteBucket method.

type S3 struct {
    client *s3.Client
}

func (s *S3) DeleteBucket(ctx context.Context, bucketName *string) error {
    input := &s3.DeleteBucketInput{
        Bucket: bucketName,
    }

    _, err := s.client.DeleteBucket(ctx, input)

    return err
}

This can be mocked up in the middleware as follows: "return output as empty" and "return with error ".

func TestS3_DeleteBucket(t *testing.T) {
    type args struct {
        ctx                context.Context
        bucketName         *string
        withAPIOptionsFunc func(*middleware.Stack) error
    }

    cases := []struct {
        name    string
        args    args
        want    error
        wantErr bool
    }{
        {
            name: "delete bucket successfully",
            args: args{
                ctx:        context.Background(),
                bucketName: aws.String("test"),
                withAPIOptionsFunc: func(stack *middleware.Stack) error {
                    return stack.Finalize.Add(
                        middleware.FinalizeMiddlewareFunc(
                            "DeleteBucketMock",
                            func(context.Context, middleware.FinalizeInput, middleware.FinalizeHandler) (middleware.FinalizeOutput, middleware.Metadata, error) {
                                return middleware.FinalizeOutput{
                                    Result: &s3.DeleteBucketOutput{},
                                }, middleware.Metadata{}, nil
                            },
                        ),
                        middleware.Before,
                    )
                },
            },
            want:    nil,
            wantErr: false,
        },
        {
            name: "delete bucket failure",
            args: args{
                ctx:        context.Background(),
                bucketName: aws.String("test"),
                withAPIOptionsFunc: func(stack *middleware.Stack) error {
                    return stack.Finalize.Add(
                        middleware.FinalizeMiddlewareFunc(
                            "DeleteBucketErrorMock",
                            func(context.Context, middleware.FinalizeInput, middleware.FinalizeHandler) (middleware.FinalizeOutput, middleware.Metadata, error) {
                                return middleware.FinalizeOutput{
                                    Result: nil,
                                }, middleware.Metadata{}, fmt.Errorf("DeleteBucketError")
                            },
                        ),
                        middleware.Before,
                    )
                },
            },
            want:    fmt.Errorf("operation error S3: DeleteBucket, DeleteBucketError"),
            wantErr: true,
        },
    }

    for _, tt := range cases {
        t.Run(tt.name, func(t *testing.T) {
            cfg, err := config.LoadDefaultConfig(
                tt.args.ctx,
                config.WithRegion("ap-northeast-1"),
                config.WithAPIOptions([]func(*middleware.Stack) error{tt.args.withAPIOptionsFunc}),
            )
            if err != nil {
                t.Fatal(err)
            }

            client := s3.NewFromConfig(cfg)
            s3Client := NewS3(client)

            err = s3Client.DeleteBucket(tt.args.ctx, tt.args.bucketName)
            if (err != nil) != tt.wantErr {
                t.Errorf("error = %#v, wantErr %#v", err, tt.wantErr)
                return
            }
            if tt.wantErr && err.Error() != tt.want.Error() {
                t.Errorf("err = %#v, want %#v", err, tt.want)
            }
        })
    }
}

Description

First of all, how to pass middleware functions to config, define a middleware function in withAPIOptionsFunc of each args, and pass LoadDefaultConfig's third argument withAPIOptionsFunc as follows in the actual test code. WithAPIOptionsFunc as the third argument of LoadDefaultConfig and set it to config.

cfg, err := config.LoadDefaultConfig(
    tt.args.ctx,
    config.WithRegion("ap-northeast-1"),
    config.WithAPIOptions([]func(*middleware.Stack) error{tt.args.withAPIOptionsFunc}),
)

Also, the following is a mock withAPIOptionsFunc that "returns output as empty".

As the name "FinalizeMiddlewareFunc" indicates, this is a middleware that modifies the response value in the Finalize step.

withAPIOptionsFunc: func(stack *middleware.Stack) error {
    return stack.Finalize.Add(
        middleware.FinalizeMiddlewareFunc(
            "DeleteBucketMock",
            func(context.Context, middleware.FinalizeInput, middleware.FinalizeHandler) (middleware.FinalizeOutput, middleware.Metadata, error) {
                return middleware.FinalizeOutput{
                    Result: &s3.DeleteBucketOutput{}, // <- This!
                }, middleware.Metadata{}, nil
            },
        ),
        middleware.Before,
    )
},

In the Result: &s3.DeleteBucketOutput{}, section, you can customize detailed output, so if you want to tweak each key/value, please change here.

If you want the API to be an "error returning" mock, just return an error instance as follows.

func(context.Context, middleware.FinalizeInput, middleware.FinalizeHandler) (middleware.FinalizeOutput, middleware.Metadata, error) {
    return middleware.FinalizeOutput{
        Result: nil,
    }, middleware.Metadata{}, fmt.Errorf("DeleteBucketError") // <- This!
},

(2) Pattern of calling multiple APIs in one test

Overall view

The next pattern is to make multiple API calls in one test case.

In this case, let's assume that the class uses CloudFormation's Client and calls DescribeStacks and DeleteStack in a single function.

type CloudFormation struct {
    client *cloudformation.Client
    waiter *cloudformation.StackDeleteCompleteWaiter
}

func (c *CloudFormation) DeleteStack(ctx context.Context, stackName *string) error {
    // ...

    output, err := c.client.DescribeStacks(ctx, input)
    if err != nil {
        return err
    }
    // ...

    if _, err := c.client.DeleteStack(ctx, input); err != nil {
        return err
    }
    // ...
}

Only withAPIOptionsFunc is described this time.

withAPIOptionsFunc: func(stack *middleware.Stack) error {
    return stack.Finalize.Add(
        middleware.FinalizeMiddlewareFunc(
            "DeleteStackOrDescribeStacksForWaiterMock",
            func(ctx context.Context, input middleware.FinalizeInput, handler middleware.FinalizeHandler) (middleware.FinalizeOutput, middleware.Metadata, error) {
                operationName := awsMiddleware.GetOperationName(ctx)
                if operationName == "DeleteStack" {
                    return middleware.FinalizeOutput{
                        Result: &cloudformation.DeleteStackOutput{},
                    }, middleware.Metadata{}, nil
                }
                if operationName == "DescribeStacks" {
                    return middleware.FinalizeOutput{
                        Result: &cloudformation.DescribeStacksOutput{
                            Stacks: []types.Stack{
                                {
                                    StackName:   aws.String("StackName"),
                                    StackStatus: "DELETE_COMPLETE",
                                },
                            },
                        },
                    }, middleware.Metadata{}, nil
                }
                return middleware.FinalizeOutput{}, middleware.Metadata{}, nil
            },
        ),
        middleware.Before,
    )
},

Description

As before, we use FinalizeMiddlewareFunc and within it the function awsMiddleware.GetOperationName(ctx).

Since the AWS SDK for Go V2 internally executes a process that puts the OperationName (name of the calling API) into the context at the time of the call, this function allows us to retrieve information about which API call it is.

This allows you to change the contents of the output returned when DeleteStack is called and when DescribeStacks is called.

operationName := awsMiddleware.GetOperationName(ctx)
if operationName == "DeleteStack" {
    return middleware.FinalizeOutput{
        Result: &cloudformation.DeleteStackOutput{},
    }, middleware.Metadata{}, nil
}
if operationName == "DescribeStacks" {
    return middleware.FinalizeOutput{
        Result: &cloudformation.DescribeStacksOutput{
            Stacks: []types.Stack{
                {
                    StackName:   aws.String("StackName"),
                    StackStatus: "DELETE_COMPLETE",
                },
            },
        },
    }, middleware.Metadata{}, nil
}

This allowed us to mock a pattern of multiple API calls in one test case with middleware.

(3) Patterns to change mock behavior depending on API arguments

Overall view

The last pattern is a pattern in which you want to change the behavior of the mock depending on the arguments (params) of the API call.

For example, in a List API, you may want to return NextMarker if no Marker is specified, and not return NextMarker if a Marker is specified.

In the List API, it is often used to return NextMarker or NextToken if the response size exceeds a certain size, and then specify it in the next request to get the next data.

In writing simple tests, you may only write tests for cases that do not return the NextMarker, but here we will use middleware that changes the behavior between cases that return the NextMarker and those that do not.

This time, in addition to the withAPIOptionsFunc, there will be one more function. In addition to Finalize, there is also an Initialize step.

Before the description, here is the big picture.

The original code is omitted this time, and the description is taken from the test file.

• New function

type markerKeyForIam struct{} 

func getNextMarkerForIamInitialize(
    ctx context.Context, in middleware.InitializeInput, next middleware.InitializeHandler,
) (
    out middleware.InitializeOutput, metadata middleware.Metadata, err error,
) {
    switch v := in.Parameters.(type) {
    case *iam.ListAttachedRolePoliciesInput:
        ctx = middleware.WithStackValue(ctx, markerKeyForIam{}, v.Marker)
    }
    return next.HandleInitialize(ctx, in)
}

• Under withAPIOptionsFunc

withAPIOptionsFunc: func(stack *middleware.Stack) error {
    err := stack.Initialize.Add(
        middleware.InitializeMiddlewareFunc(
            "GetNextMarker",
            getNextMarkerForIamInitialize, // The function I just defined
        ), middleware.Before,
    )
    if err != nil {
        return err
    }

    err = stack.Finalize.Add(
        middleware.FinalizeMiddlewareFunc(
            "ListAttachedRolePoliciesWithNextMarkerMock",
            func(ctx context.Context, input middleware.FinalizeInput, handler middleware.FinalizeHandler) (middleware.FinalizeOutput, middleware.Metadata, error) {
                marker := middleware.GetStackValue(ctx, markerKeyForIam{}).(*string)

                var nextMarker *string
                var attachedPolicies []types.AttachedPolicy
                if marker == nil {
                    nextMarker = aws.String("NextMarker") // To return NextMarker with Output
                    attachedPolicies = []types.AttachedPolicy{
                        {
                            PolicyArn:  aws.String("PolicyArn1"),
                            PolicyName: aws.String("PolicyName1"),
                        },
                    }
                    return middleware.FinalizeOutput{
                        Result: &iam.ListAttachedRolePoliciesOutput{
                            Marker:           nextMarker,
                            AttachedPolicies: attachedPolicies,
                        },
                    }, middleware.Metadata{}, nil
                } else {
                    attachedPolicies = []types.AttachedPolicy{
                        {
                            PolicyArn:  aws.String("PolicyArn2"),
                            PolicyName: aws.String("PolicyName2"),
                        },
                    }
                    return middleware.FinalizeOutput{
                        Result: &iam.ListAttachedRolePoliciesOutput{
                            Marker:           nextMarker, // Return NextMarker as nil
                            AttachedPolicies: attachedPolicies,
                        },
                    }, middleware.Metadata{}, nil
                }
            },
        ),
        middleware.Before,
    )
    return err
},

Description

Again, the pattern examples so far have used middleware with only a Finalize step, but this time the Initialize step comes into play.

The Finalize step was mainly for customizing the response, but if you want to change the behavior of the mock depending on parameters (i.e., arguments) as in this case, you need to obtain the parameters at the time of the API call, and the Initialize step is the only step that can obtain the parameters at the time of the API call. Initialize step.

Also, in order to mock the response in the Finalize step based on the parameters obtained in the Initialize step, the information obtained from the Initialize step must be passed to the next and subsequent steps.

To do this, use the function WithStackValue(ctx, key, value) from "github.com/aws/smithy-go/middleware" to store the information in the Context, and pass it to the next step in the Finalize step. GetStackValue function to retrieve the information from the Context.

Therefore, the function used for the middleware in the Initialize step is as follows. (This is defined outside of the test function, as it will be used in each test case.)

type markerKeyForIam struct{} // Used as a key for values passed to Context

func getNextMarkerForIamInitialize(
    ctx context.Context, in middleware.InitializeInput, next middleware.InitializeHandler,
) (
    out middleware.InitializeOutput, metadata middleware.Metadata, err error,
) {
    switch v := in.Parameters.(type) {
    case *iam.ListAttachedRolePoliciesInput: // At the time of a particular API call
        ctx = middleware.WithStackValue(ctx, markerKeyForIam{}, v.Marker) // Store parameter information in Context here
    }
    return next.HandleInitialize(ctx, in)
}

And withAPIOptionsFunc, wrap the above function with InitializeMiddlewareFunc and attach the middleware to the stack by stack.Initialize.Add.

withAPIOptionsFunc: func(stack *middleware.Stack) error {
    err := stack.Initialize.Add(
        middleware.InitializeMiddlewareFunc(
            "GetNextMarker",
            getNextMarkerForIamInitialize, // The function I just defined
        ), middleware.Before,
    )
    if err != nil {
        return err
    }

The parameters are then obtained at the Finalize step through the FinalizeMiddlewareFunc function, and the content of the response returned is changed based on these parameters.

    err = stack.Finalize.Add(
        middleware.FinalizeMiddlewareFunc(
            "ListAttachedRolePoliciesWithNextMarkerMock",
            func(ctx context.Context, input middleware.FinalizeInput, handler middleware.FinalizeHandler) (middleware.FinalizeOutput, middleware.Metadata, error) {

Here, we first retrieve the parameter information stored in the Context in the Initialize step.

                marker := middleware.GetStackValue(ctx, markerKeyForIam{}).(*string)

Then, based on this, we will write the logic to change the behavior of the NextMarker in the response.

As a mechanism to change the response with parameters, since there is no Marker specification in the first call (there is no way to get NextMarker since the API has not been called before), the behavior is as follows: "If there is no Marker specification, return NextMarker and make the API hit again. The second call is always a one-time call.

Since the second call always specifies the NextMarker returned in the first call as the Marker, two patterns of behavior can be realized by "not returning NextMarker if Marker is specified".

(And if the second NextMarker is not returned, it is normal for the application to terminate the process, so the API call should terminate as is.)

                var nextMarker *string
                var attachedPolicies []types.AttachedPolicy
                if marker == nil {
                    nextMarker = aws.String("NextMarker") // To return NextMarker with Output
                    attachedPolicies = []types.AttachedPolicy{
                        {
                            PolicyArn:  aws.String("PolicyArn1"),
                            PolicyName: aws.String("PolicyName1"),
                        },
                    }
                    return middleware.FinalizeOutput{
                        Result: &iam.ListAttachedRolePoliciesOutput{
                            Marker:           nextMarker,
                            AttachedPolicies: attachedPolicies,
                        },
                    }, middleware.Metadata{}, nil
                } else {
                    attachedPolicies = []types.AttachedPolicy{
                        {
                            PolicyArn:  aws.String("PolicyArn2"),
                            PolicyName: aws.String("PolicyName2"),
                        },
                    }
                    return middleware.FinalizeOutput{
                        Result: &iam.ListAttachedRolePoliciesOutput{
                            Marker:           nextMarker, // Return NextMarker as nil
                            AttachedPolicies: attachedPolicies,
                        },
                    }, middleware.Metadata{}, nil
                }

In this way, you can realize a mock that changes the behavior of the mock depending on the API arguments.

---

Finally

In other languages, you can easily mock the AWS SDK without using interfaces, but this article shows that you can do it in Go as well.

Comments

[Nicolás Parada]

Thanks. This is what I was looking for.
I was modifying the internal HTTP client to control the responses, but stopped working with some of the latest updates to the SDK. This method works great, thanks :)