Mocking AWS services with aws-smithy-mocks

Do you need mocks for AWS services to write unit tests? Why don't you give aws-smithy-mocks a shot!

Background

Suppose you are writing Rust code that interacts with AWS services through the AWS SDK for Rust, and you need to somehow mock the AWS services to test your code. There may be several options to mock AWS services, though this blog post focuses on aws-smithy-mocks and shares some practices that I have learned from using it.

What is aws-smithy-mocks?

aws-smithy-mocks is a crate for mocking AWS services for programs that use the AWS SDK for Rust. Here is an excerpt from the README.md of aws-smithy-mocks:

A flexible mocking framework for testing clients generated by smithy-rs, including all packages of the AWS SDK for Rust.

This crate provides a simple yet powerful way to mock SDK client responses for testing purposes. It uses interceptors to return stub responses, allowing you to test both happy-path and error scenarios without mocking the entire client or using traits.

The AWS documentation for Rust SDK now has a page dedicated to aws-smithy-mocks. When I first met aws-smithy-mocks, it was actually called aws-smithy-mocks-experimental. So I am happy to see that AWS is serious about the crate, at least for the time being.

Getting started with aws-smithy-mocks

Please refer to the AWS documentation for basic usage of aws-smithy-mocks. This blog post is not intended to repeat it.

Adding aws-smithy-mocks

Add aws-smithy-mocks to your development dependencies:

cargo add --dev aws-smithy-mocks

Adding the test-util feature to the AWS SDK crates

You have to enable the test-util feature for the AWS SDK crates that you want to mock with aws-smithy-mocks. Suppose you want to mock aws-sdk-dynamodb only during development, you can run the following two commands: the former for production, and the latter for development:

cargo add aws-sdk-dynamodb

cargo add --dev aws-sdk-dynamodb --features test-util

In Cargo.toml, you will find something like the following,

[dependencies]
aws-sdk-dynamodb = "1.93.0"

[dev-dependencies]
aws-sdk-dynamodb = { version = "1.93.0", features = ["test-util"] }

A drawback of duplicate dependencies is that you have to sync two versions when you want to change the SDK version.

Writing a simple unit test

Here is a simple test for an imaginary function my_function that calls the PutItem DynamoDB API:

#[tokio::test]
async fn test_dynamodb_table_put_item() {
    // 1. defines a mock behavior
    let put_item_ok = mock!(aws_sdk_dynamodb::Client::put_item)
        .then_output(|| PutItemOutput::builder().build());

    // 2. configures a client with the mock behavior
    let dynamodb = mock_client!(aws_sdk_dynamodb, [&put_item_ok]);

    // 3. executes a function to be tested with the mocked client
    let result = my_function(dynamodb).await;

    // 4. verifies the result
    assert!(result.is_ok());
}

A step-by-step explanation is below:

Defining a mock behavior.

Use the mock! macro to start mocking the behavior of an AWS service API of your choice, then chain method calls to RuleBuilder. The example mocks the PutItem DynamoDB API so that it returns a successful empty result.
```
let put_item_ok = mock!(aws_sdk_dynamodb::Client::put_item)
    .then_output(|| PutItemOutput::builder().build());
```
Configuring an SDK client with the mock behavior.

Use the mock_client! macro to configure an SDK client with mock behaviors. The example configures a DynamoDB client with the mock behavior defined at Step 1.
```
let dynamodb = mock_client!(aws_sdk_dynamodb, [&put_item_ok]);
```
Executing a function to be tested with the mocked SDK client.

How you call the function is up to your implementation, though your function needs to somehow receive an SDK client as a parameter.
```
let result = my_function(dynamodb).await;
```
Verifying the result.

Use whatever assertions you like.
```
assert!(result.is_ok());
```

Writing practical tests

I mostly use the AWS Rust SDK to write AWS Lambda functions with lambda_runtime. Techniques introduced here may also apply to tower-based systems, e.g., axum.

Parameterizing AWS SDK clients

To effectively use aws-smithy-mocks, our function to be tested has to somehow receive the AWS SDK client as a parameter.

I define a struct like the following, which retains values reused throughout the lifecycle of a Lambda instance. The derive_builder crate provides a handy derive attribute to implement the builder pattern that is especially useful during testing.

#[derive(derive_builder::Builder)]
struct SharedState {
    /// Cognito client
    cognito: aws_sdk_cognitoidentityprovider::Client,
    /// DynamoDB client
    dynamodb: aws_sdk_dynamodb::Client,
}

The main function initializes an instance of SharedState and passes it to the handler called at every Lambda invocation. It has to be wrapped in Arc so that it can be shared across async calls.

use lambda_runtime::{Error, run, service_fn};
use std::sync::Arc;

#[tokio::main]
async fn main() -> Result<(), Error> {
    let shared_state = Arc::new(SharedState::new().await?);
    run(service_fn(|req| async {
        function_handler(shared_state.clone(), req).await
    })).await
}

Isolating mocks as fixtures

Once we start writing practical mocks, our code easily gets messy. For better legibility and reusability, I extract mocks as fixtures in an isolated submodule per AWS service. I will explain it using an actual test from one of my projects.

The function being tested uses the following AWS APIs:

ListUsers (Cognito)
AdminCreateUser (Cognito)
AdminSetUserPassword (Cognito)
DeleteItem (DynamoDB)
PutItem (DynamoDB)

Messing it up

If we wrote the test without isolating mocks, we would get something like the following (non-essential details, some type definitions, and some constants are omitted for simplicity):

// lines of use declarations
use aws_sdk_cognitoidentityprovider::{
    Client as CognitoClient,
    operation::{
        admin_create_user::AdminCreateUserOutput,
        admin_set_user_password::AdminSetUserPasswordOutput,
        list_users::ListUsersOutput,
    },
    types::{AttributeType, UserType},
};
use aws_sdk_dynamodb::{
    Client as DynamodbClient,
    operation::{
        delete_item::DeleteItemOutput,
        put_item::PutItemOutput,
    },
    primitives::DateTime,
    types::AttributeValue,
};
use aws_smithy_mocks::{mock, mock_client, RuleMode};
use std::collections::HashMap;
use std::time::SystemTime;

#[tokio::test]
async fn finish_registration_of_legitimate_new_user() {
    // mocks Cognito behavior
    let list_users_empty = mock!(CognitoClient::list_users)
        .then_output(|| ListUsersOutput::builder().build());
    let admin_create_user_ok = mock!(CognitoClient::admin_create_user)
        .then_output(|| AdminCreateUserOutput::builder()
            .user(UserType::builder()
                .attributes(AttributeType::builder()
                    .name("sub")
                    .value("dummy-sub-123")
                    .build()
                    .unwrap())
                .build())
            .build());
    let admin_set_user_password_ok = mock!(CognitoClient::admin_set_user_password)
        .then_output(|| AdminSetUserPasswordOutput::builder().build());

    // mocks DynamoDB behavior
    let delete_item_session = mock!(DynamodbClient::delete_item)
        .then_output(|| {
            let ttl = DateTime::from(SystemTime::now()).secs() + 300;
            DeleteItemOutput::builder()
                .attributes("ttl", AttributeValue::N(format!("{}", ttl)))
                .attributes("state", AttributeValue::S(OK_PASSKEY_REGISTRATION.to_string()))
                .attributes("userId", AttributeValue::S("8TZ_kg_dp_pr0t7SDvGJiw".to_string()))
                .attributes("userInfo", AttributeValue::M(HashMap::from([
                    ("username".to_string(), AttributeValue::S("test".to_string())),
                    ("displayName".to_string(), AttributeValue::S("Test User".to_string())),
                ])))
                .build()
        });
    let put_item_ok = mock!(DynamodbClient::put_item)
        .then_output(|| PutItemOutput::builder().build());

    // builds SDK clients
    let cognito = mock_client!(
        aws_sdk_cognitoidentityprovider,
        RuleMode::MatchAny,
        [
            &list_users_empty,
            &admin_create_user_ok,
            &admin_set_user_password_ok,
        ]
    );
    let dynamodb = mock_client!(
        aws_sdk_dynamodb,
        RuleMode::MatchAny,
        [
            &delete_item_session,
            &put_item_ok,
        ]
    );

    // builds the state
    let shared_state = SharedStateBuilder::default()
        .webauthn(ConstantWebauthn::new(OK_PASSKEY))
        .cognito(cognito)
        .dynamodb(dynamodb)
        .build()
        .unwrap();
    let shared_state = Arc::new(shared_state);

    // verifies the function
    let res = finish_registration(
        shared_state,
        FinishRegistrationSession {
            session_id: "dummy-session-id".to_string(),
            public_key_credential: serde_json::from_str(
                OK_REGISTER_PUBLIC_KEY_CREDENTIAL,
            ).unwrap(),
        },
    ).await.unwrap();
    assert_eq!(res.user_id, "8TZ_kg_dp_pr0t7SDvGJiw");
}

In the above example, several lines are spent initializing the behavior of Cognito and DynamoDB clients. These lines may obscure the goal of the test. Imagine that there would be tens of similar tests.

There is also a minor irritation that Client conflicts between aws_sdk_cognitoidentityprovider::Client and aws_sdk_dynamodb::Client. We have to repeat fully qualified names or use aliases to disambiguate them.

Isolating mock behaviors in submodules

Isolating individual mock behaviors as fixtures will improve clarity of tests. Each fixture is a function that returns a Rule. For instance, the mock behavior of the Cognito ListUsers API in the above test can be written as the following function:

pub(crate) fn list_users_empty() -> Rule {
    mock!(Client::list_users)
        .then_output(|| ListUsersOutput::builder().build())
}

And the mock behavior of the DynamoDB DeleteItem API can be written as:

pub(crate) fn delete_item_session() -> Rule {
    mock!(Client::delete_item)
        .then_output(|| {
            let ttl = DateTime::from(SystemTime::now()).secs() + 300;
            DeleteItemOutput::builder()
                .attributes("ttl", AttributeValue::N(format!("{}", ttl)))
                .attributes("state", AttributeValue::S(OK_PASSKEY_REGISTRATION.to_string()))
                .attributes("userId", AttributeValue::S("8TZ_kg_dp_pr0t7SDvGJiw".to_string()))
                .attributes("userInfo", AttributeValue::M(HashMap::from([
                    ("username".to_string(), AttributeValue::S("test".to_string())),
                    ("displayName".to_string(), AttributeValue::S("Test User".to_string())),
                ])))
                .build()
        })
}

We will get something like the following if we group all the mock behaviors in an isolated submodule per AWS service.

pub(crate) mod mocks {
    use aws_smithy_mocks::{mock, Rule};

    pub(crate) mod cognito {
        use aws_sdk_cognitoidentityprovider::{
            Client as CognitoClient,
            operation::{
                admin_create_user::AdminCreateUserOutput,
                admin_set_user_password::AdminSetUserPasswordOutput,
                list_users::ListUsersOutput,
            },
            types::{AttributeType, UserType},
        };

        pub(crate) fn list_users_empty() -> Rule {
            mock!(Client::list_users)
                .then_output(|| ListUsersOutput::builder().build())
        }

        pub(crate) fn admin_create_user_ok() -> Rule {
            mock!(Client::admin_create_user)
                .then_output(|| AdminCreateUserOutput::builder()
                    .user(UserType::builder()
                        .attributes(AttributeType::builder()
                            .name("sub")
                            .value("dummy-sub-123")
                            .build()
                            .unwrap())
                        .build())
                    .build())
        }

        pub(crate) fn admin_set_user_password_ok() -> Rule {
            mock!(Client::admin_set_user_password)
                .then_output(|| AdminSetUserPasswordOutput::builder().build())
        }
    }

    pub(crate) mod dynamodb {
        use aws_sdk_dynamodb::{
            Client,
            operation::{
                delete_item::DeleteItemOutput,
                put_item::PutItemOutput,
            },
            primitives::DateTime,
            types::AttributeValue,
        };
        use std::collections::HashMap;
        use std::time::SystemTime;

        pub(crate) fn delete_item_session() -> Rule {
            mock!(Client::delete_item)
                .then_output(|| {
                    let ttl = DateTime::from(SystemTime::now()).secs() + 300;
                    DeleteItemOutput::builder()
                        .attributes("ttl", AttributeValue::N(format!("{}", ttl)))
                        .attributes("state", AttributeValue::S(OK_PASSKEY_REGISTRATION.to_string()))
                        .attributes("userId", AttributeValue::S("8TZ_kg_dp_pr0t7SDvGJiw".to_string()))
                        .attributes("userInfo", AttributeValue::M(HashMap::from([
                            ("username".to_string(), AttributeValue::S("test".to_string())),
                            ("displayName".to_string(), AttributeValue::S("Test User".to_string())),
                        ])))
                        .build()
                })
        }

        pub(crate) fn put_item_ok() -> Rule {
            mock!(Client::put_item)
                .then_output(|| PutItemOutput::builder().build())
        }
    }
}

As each submodule focuses on a specific AWS service, we can narrow the types to import to describe mock behaviors. There is no longer the ambiguity issue of Client.

Tidying it up

We can rewrite the above test with the mock behavior fixtures.

use aws_smithy_mocks::{mock_client, RuleMode};

#[tokio::test]
async fn finish_registration_of_legitimate_new_user() {
    // builds a Cognito client with selected mock behaviors
    let cognito = mock_client!(
        aws_sdk_cognitoidentityprovider,
        RuleMode::MatchAny,
        [
            &self::mocks::cognito::list_users_empty(),
            &self::mocks::cognito::admin_create_user_ok(),
            &self::mocks::cognito::admin_set_user_password_ok(),
        ]
    );
    // builds a DynamoDB client with selected mock behaviors
    let dynamodb = mock_client!(
        aws_sdk_dynamodb,
        RuleMode::MatchAny,
        [
            &self::mocks::dynamodb::delete_item_session(),
            &self::mocks::dynamodb::put_item_ok(),
        ]
    );

    // builds a state
    let shared_state: SharedState<ConstantWebauthnFinishRegistration> = SharedStateBuilder::default()
        .webauthn(ConstantWebauthnFinishRegistration::new(OK_PASSKEY))
        .cognito(cognito)
        .dynamodb(dynamodb)
        .build()
        .unwrap();
    let shared_state = Arc::new(shared_state);

    // verifies the function
    let res = finish_registration(
        shared_state,
        FinishRegistrationSession {
            session_id: "dummy-session-id".to_string(),
            public_key_credential: serde_json::from_str(
                OK_REGISTER_PUBLIC_KEY_CREDENTIAL,
            ).unwrap(),
        },
    ).await.unwrap();
    assert_eq!(res.user_id, "8TZ_kg_dp_pr0t7SDvGJiw");
}

Mocking error responses

Testing not only the happy path but also error scenarios is very important. In fact, most of my tests are about error scenarios.

Mocking modeled errors

A modeled error is an error defined as a variant of the enum representing the error responses from a specific AWS API, e.g., the ConditionalCheckFailedException variant of PutItemError for the error responses from the DynamoDB PutItem API.

By using the RuleBuilder::then_error method, we can make an AWS API call return a specific error response. Here is an example of making a DynamoDB PutItem API call respond with a ThrottlingException error:

pub(crate) fn put_item_throttling_exception() -> Rule {
    mock!(Client::put_item)
        .then_error(|| PutItemError::ThrottlingException(
            ThrottlingException::builder()
                .meta(ErrorMetadata::builder()
                    .code("ThrottlingException")
                    .build())
                .build(),
        ))
}

The above example configures the error code "ThrottlingException" with the meta method while building a ThrottlingException object. This is an important tip, because error handlers may depend on the error code rather than the variant type. Unfortunately, builders of error structs do not set the error code by default.

Mocking non-modeled errors

Not all of the error responses from AWS services are modeled as variants of enums. For instance, a ServiceUnavailable response is not modeled as far as I know. To make an AWS API call return a non-modeled response, we have to specify a raw HTTP response with the RuleBuilder::then_http_response method. Here is an example to let a Cognito ListUsers API call respond with a ServiceUnavailable error:

// additional use declarations
use aws_smithy_runtime_api::{
    client::orchestrator::HttpResponse,
    http::StatusCode,
};
use aws_smithy_types::body::SdkBody;

pub(crate) fn list_users_service_unavailable() -> Rule {
    mock!(Client::list_users)
        .then_http_response(|| {
            HttpResponse::new(
                StatusCode::try_from(503).unwrap(),
                SdkBody::from(r#"{"code": "ServiceUnavailable", "message": "Service unavailable."}"#),
            )
        })
}

The code field in the response JSON object determines the error code, i.e., the type = ServiceUnavailable. We also have to set the corresponding HTTP status code, 503 for a ServiceUnavailable error.

Limitation

As is often the case with mocking external services, you have to understand how a specific AWS service behaves in situations where you want to test. To overcome this limitation, you have to

Thoroughly read the AWS documentation.
Experiment with actual AWS services.

And you may also consult Generative AI.

Other options

The AWS documentation also shows other techniques to mock AWS SDK behaviors:

I have not tried them, but aws-smithy-mocks looks to cover all the aspects of the above techniques with much simpler descriptions.

Wrap up

In this blog post, I introduced,