Add HTTP mocking API

packages/smithy-aws-core/tests/functional/__init__.py

@@ -0,0 +1,2 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  SPDX-License-Identifier: Apache-2.0

packages/smithy-aws-core/tests/integration/test_user_agent_business_metrics.py → packages/smithy-aws-core/tests/functional/test_user_agent_business_metrics.py

@@ -2,9 +2,8 @@
 #  SPDX-License-Identifier: Apache-2.0
 """End-to-end User Agent 2.1 business metrics through a generated AWS client.
 
-Drives the generated awsJson1.0 client (``codegen-output/aws-json-1-0``) with a stub
-transport that records each request's User-Agent header, and checks the ``m/``
-section the SDK sends for the features each call uses.
+Drives the generated awsJson1.0 client (``codegen-output/aws-json-1-0``) and checks
+the ``m/`` section of the User-Agent the SDK sends for the features each call uses.
 """
 
 import os
@@ -24,27 +23,24 @@
 from smithy_aws_core.aio.identity.process import AsyncProcessCredentialsResolver
 from smithy_aws_core.identity.process import ProcessCredentialsResolver
 from smithy_http import HTTPResponse, tuples_to_fields
+from smithy_test import HTTPMockController, mock
+from smithy_test.aio import mock_async
 
 _JSON_CONTENT_TYPE = ("content-type", "application/x-amz-json-1.0")
 
 
+# The mock replaces the credentials resolver, so tests of credential sources
+# record requests with this transport instead.
 class _RecordingTransport:
-    """Returns canned responses and records the User-Agent of every request."""
+    """Returns empty responses and records the User-Agent of every request."""
 
     TIMEOUT_EXCEPTIONS = ()
 
-    def __init__(self, *, throttle_first: bool = False) -> None:
+    def __init__(self) -> None:
         self.user_agents: list[str] = []
-        self._throttle_first = throttle_first
 
     def send(self, request: Any, *, request_config: Any = None) -> HTTPResponse:
         self.user_agents.append(request.fields["User-Agent"].as_string())
-        if self._throttle_first and len(self.user_agents) == 1:
-            return HTTPResponse(
-                status=429,
-                fields=tuples_to_fields([_JSON_CONTENT_TYPE]),
-                body=b'{"__type":"ThrottlingException"}',
-            )
         return HTTPResponse(
             status=200,
             fields=tuples_to_fields([_JSON_CONTENT_TYPE]),
@@ -64,8 +60,8 @@ def sync(request: pytest.FixtureRequest) -> bool:
     return request.param
 
 
-def _transport(sync: bool, **kwargs: Any) -> _RecordingTransport:
-    return (_RecordingTransport if sync else _AsyncRecordingTransport)(**kwargs)
+def _transport(sync: bool) -> _RecordingTransport:
+    return (_RecordingTransport if sync else _AsyncRecordingTransport)()
 
 
 def _process_resolver(sync: bool, script: Path) -> Any:
@@ -134,52 +130,89 @@ async def _call(
     return transport.user_agents
 
 
+async def _call_mocked(
+    sync: bool,
+    tmp_path: Path,
+    metrics: str,
+    *,
+    throttle_first: bool = False,
+    **config: Any,
+) -> None:
+    missing = str(tmp_path / "missing")
+    config = {
+        "config_file_path": missing,
+        "credentials_file_path": missing,
+        "region": "us-west-2",
+        "aws_access_key_id": "AKID",
+        "aws_secret_access_key": "secret",
+        **config,
+    }
+    if sync:
+        client = JSONRPC10Client(config=JSONRPC10Config.resolve(**config))
+        with _expect_metrics(mock(client), metrics, throttle_first):
+            client.empty_input_and_empty_output(EmptyInputAndEmptyOutputInput())
+        return
+    async_client = AsyncJSONRPC10Client(
+        config=await AsyncJSONRPC10Config.resolve(**config)
+    )
+    with _expect_metrics(await mock_async(async_client), metrics, throttle_first):
+        await async_client.empty_input_and_empty_output(EmptyInputAndEmptyOutputInput())
+
+
+def _expect_metrics(
+    controller: HTTPMockController, metrics: str, throttle_first: bool
+) -> HTTPMockController:
+    expectation = controller.expect_request(
+        headers={"User-Agent": lambda ua: _business_metrics(ua) == [metrics]}
+    )
+    if throttle_first:
+        expectation.times(2).respond_with(
+            status_code=429,
+            headers=dict([_JSON_CONTENT_TYPE]),
+            body='{"__type":"ThrottlingException"}',
+        )
+    expectation.respond_with(headers=dict([_JSON_CONTENT_TYPE]), body="{}")
+    return controller
+
+
 async def test_reports_standard_retry_mode_by_default(
     sync: bool, tmp_path: Path
 ) -> None:
-    user_agents = await _call(sync, tmp_path, _transport(sync))
-    assert _business_metrics(user_agents[0]) == ["m/E,e"]
+    await _call_mocked(sync, tmp_path, "m/E,e")
 
 
 async def test_reports_adaptive_retry_mode(sync: bool, tmp_path: Path) -> None:
-    user_agents = await _call(sync, tmp_path, _transport(sync), retry_mode="adaptive")
-    assert _business_metrics(user_agents[0]) == ["m/F,e"]
+    await _call_mocked(sync, tmp_path, "m/F,e", retry_mode="adaptive")
 
 
 async def test_reports_adaptive_retry_mode_from_environment(
     sync: bool, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
 ) -> None:
     monkeypatch.setenv("AWS_RETRY_MODE", "adaptive")
-    user_agents = await _call(sync, tmp_path, _transport(sync))
-    assert _business_metrics(user_agents[0]) == ["m/F,e"]
+    await _call_mocked(sync, tmp_path, "m/F,e")
 
 
 async def test_reports_endpoint_override(sync: bool, tmp_path: Path) -> None:
-    user_agents = await _call(
-        sync, tmp_path, _transport(sync), endpoint_uri="https://localhost:8443"
-    )
-    assert _business_metrics(user_agents[0]) == ["m/E,N,e"]
+    await _call_mocked(sync, tmp_path, "m/E,N,e", endpoint_uri="https://localhost:8443")
 
 
 async def test_reports_endpoint_override_from_environment(
     sync: bool, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
 ) -> None:
     monkeypatch.setenv("AWS_ENDPOINT_URL_JSON_RPC_10", "https://localhost:8443")
-    user_agents = await _call(sync, tmp_path, _transport(sync))
-    assert _business_metrics(user_agents[0]) == ["m/E,N,e"]
+    await _call_mocked(sync, tmp_path, "m/E,N,e")
 
 
 async def test_retried_call_sends_same_business_metrics_once(
     sync: bool, tmp_path: Path
 ) -> None:
-    user_agents = await _call(
+    await _call_mocked(
         sync,
         tmp_path,
-        _transport(sync, throttle_first=True),
+        "m/E,N,e",
+        throttle_first=True,
         endpoint_uri="https://localhost:8443",
     )
-    assert len(user_agents) == 2
-    assert [_business_metrics(ua) for ua in user_agents] == [["m/E,N,e"], ["m/E,N,e"]]
 
 
 async def test_reports_process_resolver_used_directly(

packages/smithy-core/tests/functional/test_retries.py

@@ -49,6 +49,8 @@ def retry_operation(
         if isinstance(response, Exception):
             error = response
         else:
+            # TODO: Clients don't retry 5xx responses: the protocols mark unmodeled
+            # errors retry-safe only for 408 and 429. This sets it so 5xx retries.
             error = CallError(
                 fault="server" if response >= 500 else "client",
                 message=f"HTTP {response}",

packages/smithy-http/.changes/next-release/smithy-http-enhancement-76678ead0b6642e5800a99d5a245c8a8.json

@@ -0,0 +1,4 @@
+{
+  "type": "enhancement",
+  "description": "Allow MockHTTPClient and AsyncMockHTTPClient responses to have stream bodies."
+}

packages/smithy-http/src/smithy_http/aio/testing.py

@@ -38,4 +38,6 @@ async def send(
         :raises MockHTTPClientError: If no responses are queued.
         """
         status, fields, body = self._pop_response(request)
-        return HTTPResponse(status=status, fields=fields, body=list_async([body]))
+        if isinstance(body, bytes | bytearray):
+            body = list_async([bytes(body)])
+        return HTTPResponse(status=status, fields=fields, body=body)

packages/smithy-http/src/smithy_http/testing/mockhttp.py

@@ -6,6 +6,8 @@
 from io import BytesIO
 from typing import Any, Self
 
+from smithy_core.interfaces import AnyStreamingBlob
+
 from smithy_http import HTTPResponse, tuples_to_fields
 from smithy_http.interfaces import (
     Fields,
@@ -37,7 +39,7 @@ def __init__(
         client.
         """
         self._client_config = client_config
-        self._responses: deque[tuple[int, Fields, bytes]] = deque()
+        self._responses: deque[tuple[int, Fields, AnyStreamingBlob]] = deque()
         self._captured_requests: list[HTTPRequest] = []
 
     def __deepcopy__(self, memo: Any) -> Self:
@@ -57,17 +59,19 @@ def add_response(
         self,
         status: int = 200,
         headers: list[tuple[str, str]] | None = None,
-        body: bytes = b"",
+        body: AnyStreamingBlob = b"",
     ) -> None:
         """Queue a response for the next request.
 
         :param status: HTTP status code.
         :param headers: HTTP response headers as list of (name, value) tuples.
-        :param body: Response body as bytes.
+        :param body: Response body as bytes, or a stream the client can read.
         """
         self._responses.append((status, tuples_to_fields(headers or []), body))
 
-    def _pop_response(self, request: HTTPRequest) -> tuple[int, Fields, bytes]:
+    def _pop_response(
+        self, request: HTTPRequest
+    ) -> tuple[int, Fields, AnyStreamingBlob]:
         """Capture the request and return the next queued status, fields and body.
 
         :raises MockHTTPClientError: If no responses are queued.
@@ -101,7 +105,9 @@ def send(
         :raises MockHTTPClientError: If no responses are queued.
         """
         status, fields, body = self._pop_response(request)
-        return HTTPResponse(status=status, fields=fields, body=BytesIO(body))
+        if isinstance(body, bytes | bytearray):
+            body = BytesIO(body)
+        return HTTPResponse(status=status, fields=fields, body=body)
 
 
 class MockHTTPClientError(Exception):

packages/smithy-http/tests/functional/__init__.py

@@ -0,0 +1,2 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  SPDX-License-Identifier: Apache-2.0

packages/smithy-http/tests/functional/test_streaming_payloads.py

@@ -0,0 +1,73 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  SPDX-License-Identifier: Apache-2.0
+from io import BytesIO
+
+import pytest
+
+pytest.importorskip("restjson")
+
+from restjson.aio.client import AsyncRestJsonProtocolClient
+from restjson.aio.config import AsyncRestJsonProtocolConfig
+from restjson.client import RestJsonProtocolClient
+from restjson.config import RestJsonProtocolConfig
+from restjson.models import StreamingTraitsRequireLengthInput
+from smithy_core.aio.types import AsyncBytesReader
+from smithy_core.exceptions import SerializationError
+from smithy_core.interfaces import StreamingBlob
+from smithy_test import HTTPMockController, HTTPMockExpectation, mock
+from smithy_test.aio import mock_async
+
+PAYLOAD = b'{"hello":"world"}'
+
+
+def test_streaming_payloads_are_sent_the_same() -> None:
+    client = RestJsonProtocolClient(RestJsonProtocolConfig.resolve(region="us-west-2"))
+    with mock(client) as controller:
+        _expect_payload(controller).times(3)
+
+        for blob in _blobs():
+            client.streaming_traits_require_length(
+                StreamingTraitsRequireLengthInput(blob=blob)
+            )
+
+
+async def test_async_streaming_payloads_are_sent_the_same() -> None:
+    config = await AsyncRestJsonProtocolConfig.resolve(region="us-west-2")
+    async with AsyncRestJsonProtocolClient(config) as client:
+        with await mock_async(client) as controller:
+            _expect_payload(controller).times(3)
+
+            for blob in _blobs():
+                await client.streaming_traits_require_length(
+                    StreamingTraitsRequireLengthInput(blob=blob)
+                )
+
+
+async def test_unsized_streams_are_rejected_before_sending() -> None:
+    config = await AsyncRestJsonProtocolConfig.resolve(region="us-west-2")
+    async with AsyncRestJsonProtocolClient(config) as client:
+        controller = await mock_async(client)
+        _expect_payload(controller)
+
+        with pytest.raises(SerializationError, match="content length"):
+            await client.streaming_traits_require_length(
+                StreamingTraitsRequireLengthInput(blob=AsyncBytesReader(PAYLOAD))
+            )
+
+    assert controller.requests == []
+
+
+def _blobs() -> list[StreamingBlob]:
+    return [PAYLOAD, bytearray(PAYLOAD), BytesIO(PAYLOAD)]
+
+
+def _expect_payload(controller: HTTPMockController) -> HTTPMockExpectation:
+    return controller.expect_request(
+        method="POST",
+        path="/StreamingTraitsRequireLength",
+        headers={
+            "Content-Type": "application/octet-stream",
+            "Content-Length": str(len(PAYLOAD)),
+        },
+        body=PAYLOAD.decode(),
+    )

packages/smithy-http/tests/unit/aio/testing/test_mockhttp.py

@@ -2,6 +2,7 @@
 #  SPDX-License-Identifier: Apache-2.0
 
 import pytest
+from smithy_core.aio.types import AsyncBytesReader
 from smithy_core.aio.utils import read_streaming_blob_async
 from smithy_http.aio.testing import AsyncMockHTTPClient
 from smithy_http.testing import MockHTTPClientError, create_test_request
@@ -121,3 +122,14 @@ async def test_captured_requests_copy():
     assert captured1 is not captured2
     # But with same content
     assert len(captured1) == len(captured2) == 1
+
+
+async def test_responds_with_a_stream():
+    client = AsyncMockHTTPClient()
+    stream = AsyncBytesReader(b"streamed")
+    client.add_response(body=stream)
+
+    response = await client.send(create_test_request())
+
+    assert response.body is stream
+    assert await read_streaming_blob_async(response.body) == b"streamed"

packages/smithy-http/tests/unit/testing/test_mockhttp.py

@@ -1,6 +1,8 @@
 #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #  SPDX-License-Identifier: Apache-2.0
 
+from io import BytesIO
+
 import pytest
 from smithy_core.utils import read_streaming_blob
 from smithy_http.testing import MockHTTPClient, MockHTTPClientError, create_test_request
@@ -27,3 +29,14 @@ def test_queued_responses_fifo():
 def test_raises_without_queued_response():
     with pytest.raises(MockHTTPClientError, match="No responses queued"):
         MockHTTPClient().send(create_test_request())
+
+
+def test_responds_with_a_stream():
+    client = MockHTTPClient()
+    stream = BytesIO(b"streamed")
+    client.add_response(body=stream)
+
+    response = client.send(create_test_request())
+
+    assert response.body is stream
+    assert read_streaming_blob(response.body) == b"streamed"

packages/smithy-test/.changes/next-release/smithy-test-feature-374bd727a36340568563fe5d5fd5eb36.json

@@ -0,0 +1,4 @@
+{
+  "type": "feature",
+  "description": "Add HTTP mocking API"
+}

packages/smithy-test/pyproject.toml

@@ -29,6 +29,8 @@ classifiers = [
 ]
 dependencies = [
     "smithy-core~=0.10.0",
+    "smithy-http~=0.7.0",
+    "smithy-aws-core~=0.13.0",
 ]
 
 [project.urls]