Moto + aiobotocore

aws.py

from collections.abc import Awaitable, Callable, Iterator
from dataclasses import dataclass
from types import TracebackType
from typing import Any, TypeVar

import aiobotocore
import aiobotocore.endpoint
import aiobotocore.httpchecksum
import aiobotocore.response
import botocore.awsrequest


_T = TypeVar("_T")
_T2 = TypeVar("_T2")


@dataclass
class _PatchedAWSResponseContent:
    """Patched version of `botocore.awsrequest.AWSResponse.content`"""

    content: bytes | Awaitable[bytes]

    def __await__(self) -> Iterator[bytes]:
        async def _generate_async() -> bytes:
            if isinstance(self.content, Awaitable):
                return await self.content
            else:
                return self.content

        return _generate_async().__await__()

    def decode(self, encoding: str) -> str:
        assert isinstance(self.content, bytes)
        return self.content.decode(encoding)


class PatchedAWSResponse:
    """Patched version of `botocore.awsrequest.AWSResponse`"""

    def __init__(self, response: botocore.awsrequest.AWSResponse) -> None:
        self._response = response
        self.status_code = response.status_code
        self.content = _PatchedAWSResponseContent(response.content)
        self.raw = response.raw
        self.headers = response.headers
        if not hasattr(self.raw, "raw_headers"):
            self.raw.raw_headers = {}


def _convert_to_response_dict_factory(
    original: Callable[[botocore.awsrequest.AWSResponse, _T], Awaitable[_T2]],
) -> Callable[[botocore.awsrequest.AWSResponse, _T], Awaitable[_T2]]:
    """Factory for patching `aiobotocore.endpoint.convert_to_response_dict`"""

    async def convert_to_response_dict(http_response: botocore.awsrequest.AWSResponse, operation_model: _T) -> _T2:
        return await original(PatchedAWSResponse(http_response), operation_model)  # type: ignore[arg-type]

    return convert_to_response_dict


aiobotocore.endpoint.convert_to_response_dict = _convert_to_response_dict_factory(  # type: ignore[assignment]
    aiobotocore.endpoint.convert_to_response_dict  # type: ignore[arg-type]
)


async def _patched_read(self: aiobotocore.response.StreamingBody, amt: Any = None) -> Any:  # noqa: ARG001
    """Patched version of `aiobotocore.response.StreamingBody.read`"""
    # Get the full content.
    return self.__wrapped__.getvalue()


aiobotocore.response.StreamingBody.read = _patched_read  # type: ignore[method-assign]


async def _patched_aenter(self: aiobotocore.response.StreamingBody) -> Any:
    """Patched version of `aiobotocore.response.StreamingBody.__aenter__`"""
    result = self.__wrapped__.__enter__()

    class _AsyncStreamResponse:
        async def read(self) -> bytes:
            # Get full content.
            return result.getvalue()  # type: ignore[no-any-return]

    return _AsyncStreamResponse()


aiobotocore.response.StreamingBody.__aenter__ = _patched_aenter  # type: ignore[assignment, method-assign]


async def _patched_aexit(
    self: aiobotocore.response.StreamingBody,
    exc_type: type[BaseException] | None,
    exc_val: BaseException | None,
    exc_tb: TracebackType | None,
) -> Any:
    """Patched version of `aiobotocore.response.StreamingBody.__aexit__`"""
    return self.__wrapped__.__exit__(exc_type, exc_val, exc_tb)


aiobotocore.response.StreamingBody.__aexit__ = _patched_aexit  # type: ignore[assignment, method-assign]

# Remove async versions of functions which are not compatible with Moto.
del aiobotocore.httpchecksum.AioAwsChunkedWrapper._make_chunk  # noqa: SLF001
del aiobotocore.httpchecksum.AioAwsChunkedWrapper.read
del aiobotocore.httpchecksum.StreamingChecksumBody.read

versions.txt

aiobotocore                           2.23.2
moto                                  5.1.9
mypy                                  1.17.0
pyright                               1.1.403
ruff                                  0.11.7

Been posting in this thread for a few years now - aio-libs/aiobotocore#755, decided to put my most up-to-date code in a gist.

Works with these packages:

moto (5.1.0)
aiobotocore (2.20.0)
mypy (1.15.0)
pyright (1.1.394)
ruff (0.9.7)

Thanks for the updates! It seems to be broken with aiobotocore > 2.23.0 again unfortunately.

Updated the gist, works for these versions:

aiobotocore                           2.23.2
moto                                  5.1.9
mypy                                  1.17.0
pyright                               1.1.403
ruff                                  0.11.7