Advanced client stubbing in the AWS SDK for Ruby Version 3

The AWS SDK for Ruby provides a robust set of features for stubbing your clients, to make unit tests easier and less fragile. Many of you have used these features to stub your client calls. But in this post, we’re going to explore both a new stub feature in version 3 of the AWS SDK for Ruby, and some advanced stubbing techniques you can add to your testing toolbox.

How to stub a client

The Aws::ClientStubs documentation has several examples of how to write client stubs, but in the most basic form, you do the following:

Set the :stub_responses parameter at client creation.
Write out the stubbed responses you want, sequentially.

For a simple example of stubbing an operation, consider the following.

s3 = Aws::S3::Client.new(stub_responses: true)
s3.stub_responses(:list_buckets, { buckets: [{ name: "my-bucket" }] })
s3.list_buckets.buckets.map(&:name) #=> ['my-bucket']

You can also stub the same operation to provide different responses for sequential calls.

s3 = Aws::S3::Client.new(stub_responses: true)
s3.stub_responses(:get_object,
  'NoSuchKey',
  { body: "Hello!" }
)
begin
  s3.get_object(bucket: "test", key: "test")
rescue Aws::S3::Errors::NoSuchKey
  # You don't NEED to do this, but maybe your test function would be doing this.
  s3.put_object(bucket: "test", key: "test", body: "Doesn't matter")
end
s3.get_object(bucket: "test", key: "test").body.read #=> "Hello!"

This works pretty well for most test use cases, but it can be fragile in others. We’re expecting that API calls will come in a certain sequence, and for that #put_object call, the :body parameter value didn’t matter at all – the stub is fixed. In some cases, we want our stubs to have a bit of dynamic logic, and that’s where dynamic client stubbing is an option.

Dynamic client stubbing

The #stub_responses method accepts more than static response objects in sequence. You can also provide a Proc object, which is able to inspect the request context and dynamically determine a response. To take the previous Amazon S3 example, we could have an in-memory bucket that dynamically tracks objects in the database, and can even be pre-seeded.

bucket = {
  "SeededKey" => { body: "Hello!" }
}
s3 = Aws::S3::Client.new(stub_responses: true)
s3.stub_responses(:get_object, -> (context) {
  obj = bucket[context.params[:key]]
  if obj
    obj
  else
    'NoSuchKey'
  end
})
s3.stub_responses(:put_object, -> (context) {
  bucket[context.params[:key]] = { body: context.params[:body] }
  {}
})
begin
  s3.get_object(bucket: "test", key: "test")
rescue Aws::S3::Errors::NoSuchKey
  s3.put_object(bucket: "test", key: "test", body: "DOES matter!")
end
s3.get_object(bucket: "test", key: "test").body.read #=> "DOES matter!"
s3.get_object(bucket: "test", key: "SeededKey").body.read #=> "Hello!"

We’ll take this even further in the final example.

New feature: recorded API requests

While developing the aws-record gem, we discovered that we needed additional testing functionality around client calls. When creating Amazon DynamoDB tables from attribute and metadata specifications, the main thing we wanted to test was that the #create_table parameters exactly matched what we would have otherwise handcrafted. The stubbed response was mostly irrelevant.

To solve that problem, our tests added a lightweight “handler” that recorded the request parameters. These could then be inspected by our tests and compared to expectations. Starting with version 3.23.0 of aws-sdk-core, this is now a built-in feature of stubbed SDK clients!

You can access this set of API requests directly from your stubbed client, as follows.

require 'aws-sdk-s3'

s3 = Aws::S3::Client.new(stub_responses: true)
s3.create_bucket(bucket: "foo")
s3.api_requests.size # => 1
create_bucket_request = s3.api_requests.first
create_bucket_request[:params][:bucket] # => "foo"

To see how this is used, here is how we’d rewrite one of the Aws::Record::TableConfig unit tests and its prerequisites.

require 'rspec'
require 'aws-record'

class TestModel
  include Aws::Record

  string_attr :hk, hash_key: true
  string_attr :rk, range_key: true
end

module Aws
  module Record
    describe TableConfig do
      describe "#migrate!" do
        it 'will attempt to create the remote table if it does not exist' do
          cfg = TableConfig.define do |t|
            t.model_class(TestModel)
            t.read_capacity_units(1)
            t.write_capacity_units(1)
            t.client_options(stub_responses: true)
          end
          stub_client = cfg.client

          # Sequential responses are appropriate, so a proc stub is more than
          # we need.
          stub_client.stub_responses(
            :describe_table,
            'ResourceNotFoundException',
            { table: { table_status: "ACTIVE" } }
          )
          cfg.migrate!

          # We don't need to know how many calls happened around describing
          # state. We can find the call we care about.
          create_table_request = stub_client.api_requests.find do |req|
            req[:operation_name] == :create_table
          end

          # Parameters are separated into their own key. Full context is available
          # if you want it.
          expect(create_table_request[:params]).to eq(
            table_name: "TestModel",
            provisioned_throughput:
            {
              read_capacity_units: 1,
              write_capacity_units: 1
            },
            key_schema: [
              {
                attribute_name: "hk",
                key_type: "HASH"
              },
              {
                attribute_name: "rk",
                key_type: "RANGE"
              }
            ],
            attribute_definitions: [
              {
                attribute_name: "hk",
                attribute_type: "S"
              },
              {
                attribute_name: "rk",
                attribute_type: "S"
              }
            ]
          )
        end
      end
    end
  end
end

This is one way to make tests a little less fragile, and test both how you handle client responses (via stubbing) and how you form your requests (via the #api_requests log of requests made to the client).

Advanced stubbing test example

Let’s bring all of this together into a runnable test file.

Let’s say we’re testing a class that interacts with Amazon S3. We’re performing relatively basic operations around writing and retrieving objects, but don’t want to keep track of which fixed stubs go in which order.

The code below creates a very simple “Fake S3” with an in-memory hash, and implements the #create_bucket, #get_object, and #put_object APIs for their basic parameters. With this, we’re able to verify that our code makes the client calls we intend and handles the responses, without tracking API client call order details as the tests get more complex.

(A caveat: In reality, you likely wouldn’t be testing the client stubbing functionality directly. Instead, you’d be calling into your own functions and then making these checks. However, for example purposes, the file is standalone.)

require 'rspec'
require 'aws-sdk-s3'

describe "Enhanced Stubbing Example Tests" do
  let(:fake_s3) { {} }
  let(:client) do
    client = Aws::S3::Client.new(stub_responses: true)
    client.stub_responses(
      :create_bucket, ->(context) {
        name = context.params[:bucket]
        if fake_s3[name]
          'BucketAlreadyExists' # standalone strings are treated as exceptions
        else
          fake_s3[name] = {}
          {}
        end
      }
    )
    client.stub_responses(
      :get_object, ->(context) {
        bucket = context.params[:bucket]
        key = context.params[:key]
        b_contents = fake_s3[bucket]
        if b_contents
          obj = b_contents[key]
          if obj
            { body: obj }
          else
            'NoSuchKey'
          end
        else
          'NoSuchBucket'
        end
      }
    )
    client.stub_responses(
      :put_object, ->(context) {
        bucket = context.params[:bucket]
        key = context.params[:key]
        body = context.params[:body]
        b_contents = fake_s3[bucket]
        if b_contents
          b_contents[key] = body
          {}
        else
          'NoSuchBucket'
        end
      }
    )
    client
  end

  it "raises an exception when a bucket is created twice" do
    client.create_bucket(bucket: "foo")
    client.create_bucket(bucket: "bar")
    expect {
      client.create_bucket(bucket: "foo")
    }.to raise_error(Aws::S3::Errors::BucketAlreadyExists)
    expect(client.api_requests.size).to eq(3)
  end

  context "bucket operations" do
    before do
      client.create_bucket(bucket: "test")
    end

    it "can write and retrieve an object" do
      client.put_object(bucket: "test", key: "obj", body: "Hello!")
      obj = client.get_object(bucket: "test", key: "obj")
      expect(obj.body.read).to eq("Hello!")
      expect(client.api_requests.size).to eq(3)
      expect(client.api_requests.last[:params]).to eq(
        bucket: "test",
        key: "obj"
      )
    end

    it "raises the appropriate exception when a bucket doesn't exist" do
      expect {
        client.put_object(
          bucket: "sirnotappearinginthistest",
          key: "knight_sayings",
          body: "Ni!"
        )
      }.to raise_error(Aws::S3::Errors::NoSuchBucket)
      expect(client.api_requests.size).to eq(2)
    end

    it "raises the appropriate exception when an object doesn't exist" do
      expect {
        client.get_object(bucket: "test", key: "404NoSuchKey")
      }.to raise_error(Aws::S3::Errors::NoSuchKey)
      expect(client.api_requests.size).to eq(2)
    end
  end
end

Conclusion

Client stubs in the AWS SDK for Ruby are a powerful tool for unit testing. They provide a way to test without hitting the network, but allow your code to behave like it’s calling the AWS API clients without having to form mocks for full response objects. This can bring your testing closer to “the real thing” and help you develop code with the SDK with increased confidence.

AWS Developer Tools Blog