How to stub S3.Object.wait_until

I have been tasked with writing tests for an s3 uploading function which uses S3.Object.wait_until_exists to wait for upload to complete and get the content length of the upload to return it.

But so far I am failing to stub head_object for the waiter.

I have explored and found the waiter has two acceptors:

if HTTP code == 200, accept
if HTTP code == 404, retry

I don't know how to explain in text more so instead here is an MRE.

from datetime import datetime
from io import BytesIO

import boto3
import botocore
import botocore.stub

testing_bucket = "bucket"
testing_key = "key/of/object"
testing_data = b"data"

s3 = boto3.resource("s3")


def put():
    try:
        o = s3.Object(testing_bucket, testing_key)
        o.load()  # head_object * 1
    except botocore.exceptions.ClientError as e:
        if e.response["Error"]["Code"] == "NoSuchKey":
            etag = ""
        else:
            raise e
    else:
        etag = o.e_tag
    try:
        o.upload_fileobj(BytesIO(testing_data))  # put_object * 1
    except botocore.exceptions.ClientError as e:
        raise e
    else:
        o.wait_until_exists(IfNoneMatch=etag)  # head_object * n until accepted
        return o.content_length  # not sure if calling head_object again


with botocore.stub.Stubber(s3.meta.client) as s3_stub:
    s3_stub.add_response(
        method="head_object",
        service_response={
            "ETag": "fffffffe",
            "ContentLength": 0,
        },
        expected_params={
            "Bucket": testing_bucket,
            "Key": testing_key,
        },
    )
    s3_stub.add_response(
        method="put_object",
        service_response={},
        expected_params={
            "Bucket": testing_bucket,
            "Key": testing_key,
            "Body": botocore.stub.ANY,
        },
    )
    s3_stub.add_response(  # cause time to increase by 5 seconds per response
        method="head_object",
        service_response={
            "ETag": "ffffffff",
            "AcceptRanges": "bytes",
            "ContentLength": len(testing_data),
            "LastModified": datetime.now(),
            "Metadata": {},
            "VersionId": "null",
        },
        expected_params={
            "Bucket": testing_bucket,
            "Key": testing_key,
            "IfNoneMatch": "fffffffe",
        },
    )
    print(put())  # should print 4

And running the above gives:

time python mre.py
Traceback (most recent call last):
  File "/tmp/mre.py", line 72, in <module>
    put()
  File "/tmp/mre.py", line 30, in put
    o.wait_until_exists(IfNoneMatch=etag)  # head_object * 1
  File "/tmp/.tox/py310/lib/python3.10/site-packages/boto3/resources/factory.py", line 413, in do_waiter
    waiter(self, *args, **kwargs)
  File "/tmp/.tox/py310/lib/python3.10/site-packages/boto3/resources/action.py", line 215, in __call__
    response = waiter.wait(**params)
  File "/tmp/.tox/py310/lib/python3.10/site-packages/botocore/waiter.py", line 55, in wait
    Waiter.wait(self, **kwargs)
  File "/tmp/.tox/py310/lib/python3.10/site-packages/botocore/waiter.py", line 343, in wait
    response = self._operation_method(**kwargs)
  File "/tmp/.tox/py310/lib/python3.10/site-packages/botocore/waiter.py", line 93, in __call__
    return self._client_method(**kwargs)
  File "/tmp/.tox/py310/lib/python3.10/site-packages/botocore/client.py", line 508, in _api_call
    return self._make_api_call(operation_name, kwargs)
  File "/tmp/.tox/py310/lib/python3.10/site-packages/botocore/client.py", line 878, in _make_api_call
    request_dict = self._convert_to_request_dict(
  File "/tmp/.tox/py310/lib/python3.10/site-packages/botocore/client.py", line 936, in _convert_to_request_dict
    api_params = self._emit_api_params(
  File "/tmp/.tox/py310/lib/python3.10/site-packages/botocore/client.py", line 969, in _emit_api_params
    self.meta.events.emit(
  File "/tmp/.tox/py310/lib/python3.10/site-packages/botocore/hooks.py", line 412, in emit
    return self._emitter.emit(aliased_event_name, **kwargs)
  File "/tmp/.tox/py310/lib/python3.10/site-packages/botocore/hooks.py", line 256, in emit
    return self._emit(event_name, kwargs)
  File "/tmp/.tox/py310/lib/python3.10/site-packages/botocore/hooks.py", line 239, in _emit
    response = handler(**kwargs)
  File "/tmp/.tox/py310/lib/python3.10/site-packages/botocore/stub.py", line 376, in _assert_expected_params
    self._assert_expected_call_order(model, params)
  File "/tmp/.tox/py310/lib/python3.10/site-packages/botocore/stub.py", line 352, in _assert_expected_call_order
    raise UnStubbedResponseError(
botocore.exceptions.UnStubbedResponseError: Error getting response stub for operation HeadObject: Unexpected API Call: A call was made but no additional calls expected. Either the API Call was not stubbed or it was called multiple times.
python mre.py  0.39s user 0.19s system 9% cpu 5.859 total

Or with 2 answer, same thing with python mre.py 0.40s user 0.20s system 5% cpu 10.742 total.

CodePudding user response：

I found a solution for this, as highlighted the waiter waits for a 200 status code, adding it to the response like the following works:

s3_stub.add_response(
        method="head_object",
        service_response={
            "ETag": "ffffffff",
            "AcceptRanges": "bytes",
            "ContentLength": len(testing_data),
            "LastModified": datetime.now(),
            "Metadata": {},
            "VersionId": "null",
            "ResponseMetadata": {"HTTPStatusCode": 200},
        },
        expected_params={
            "Bucket": testing_bucket,
            "Key": testing_key,
            "IfNoneMatch": "fffffffe",
        },
    )