diff --git a/packages/smithy-core/.changes/next-release/smithy-core-enhancement-ce57b5cfed2740908d9e5222d3f58f47.json b/packages/smithy-core/.changes/next-release/smithy-core-enhancement-ce57b5cfed2740908d9e5222d3f58f47.json new file mode 100644 index 000000000..f5c7be6c4 --- /dev/null +++ b/packages/smithy-core/.changes/next-release/smithy-core-enhancement-ce57b5cfed2740908d9e5222d3f58f47.json @@ -0,0 +1,4 @@ +{ + "type": "enhancement", + "description": "Added retries for requests with a seekable streaming payload, which are resent from the position the stream was at when the operation was called." +} \ No newline at end of file diff --git a/packages/smithy-core/src/smithy_core/aio/client.py b/packages/smithy-core/src/smithy_core/aio/client.py index ec43db2be..f2e1eb7a2 100644 --- a/packages/smithy-core/src/smithy_core/aio/client.py +++ b/packages/smithy-core/src/smithy_core/aio/client.py @@ -30,7 +30,7 @@ from ..interfaces.auth import AuthOption, AuthSchemeResolver from ..schemas import APIOperation from ..serializers import SerializeableShape -from ..shapes import ShapeID +from ..shapes import ShapeID, ShapeType from ..types import PropertyKey from ..types import TypedProperties as _TypedProperties from .eventstream import DuplexEventStream, InputEventStream, OutputEventStream @@ -92,8 +92,26 @@ class ClientCall[I: SerializeableShape, O: DeserializeableShape]: """The retry scope for the operation.""" def retryable(self) -> bool: - # TODO: check to see if the stream is seekable - return self.operation.input_stream_member is None + """Whether the operation can be retried. + + Operations with an input event stream are never retried. A request for any + other operation is retried only if its body can be rewound to where it started. + """ + member = self.operation.input_stream_member + return member is None or member.shape_type is not ShapeType.UNION + + +async def _rewind_position(body: Any) -> int | None: + """Find where a request body is rewound to before it's resent. + + :returns: The body's current position, or None if it can't be rewound. + """ + if body is None or isinstance(body, bytes | bytearray): + return 0 + try: + return await seek(body, 0, 1) + except (OSError, ValueError, NotImplementedError): + return None class RequestPipeline[TRequest: Request, TResponse: Response]: @@ -382,7 +400,8 @@ async def _retry[I: SerializeableShape, O: DeserializeableShape]( request_context: RequestContext[I, TRequest], request_future: Future[RequestContext[I, TRequest]] | None, ) -> OutputContext[I, O, TRequest | None, TResponse | None]: - if not call.retryable(): + body = request_context.transport_request.body + if not call.retryable() or (start := await _rewind_position(body)) is None: return await self._send_with_rate_limiting( call, request_context, request_future ) @@ -444,7 +463,11 @@ async def _retry[I: SerializeableShape, O: DeserializeableShape]( retry_token.retry_delay, ) - await seek(request_context.transport_request.body, 0) + try: + await seek(body, start) + except (OSError, ValueError, NotImplementedError) as e: + output_context.response.add_note(f"Retry not attempted: {e}") + return output_context else: await retry_strategy.record_success(token=retry_token) return output_context diff --git a/packages/smithy-core/tests/unit/aio/_pipeline_harness.py b/packages/smithy-core/tests/unit/aio/_pipeline_harness.py index c4420181a..1f3f18a23 100644 --- a/packages/smithy-core/tests/unit/aio/_pipeline_harness.py +++ b/packages/smithy-core/tests/unit/aio/_pipeline_harness.py @@ -12,6 +12,7 @@ from smithy_core.documents import TypeRegistry from smithy_core.endpoints import EndpointResolverParams from smithy_core.interceptors import InterceptorChain +from smithy_core.prelude import BLOB from smithy_core.response import EMPTY_RESPONSE_METADATA, ResponseMetadata from smithy_core.schemas import APIOperation, Schema from smithy_core.serializers import ShapeSerializer @@ -86,7 +87,29 @@ def deserialize(cls, deserializer: ShapeDeserializer) -> Self: effective_auth_schemes=[], ) -# ``ClientCall.retryable()`` is False for operations with a streaming input, so +_DATA_STREAM_INPUT_SCHEMA = Schema.collection( + id=ShapeID("com.example#DataStreamInput"), + members={"body": {"target": BLOB, "traits": [StreamingTrait()]}}, +) + +# An operation whose input is a data stream rather than an event stream. +DATA_STREAM_OPERATION = APIOperation( + input=StubInput, + output=StubOutput, + static_schema=( + ShapeType.OPERATION.value, + "com.example", + "DataStreamOperation", + 0, + _DATA_STREAM_INPUT_SCHEMA, + _UNARY_OUTPUT_SCHEMA, + (), + ), + error_registry=TypeRegistry({}), + effective_auth_schemes=[], +) + +# ``ClientCall.retryable()`` is False for operations with an input event stream, so # OPERATION above never enters the retry loop. This one has no input stream and so # exercises it. UNARY_OPERATION = APIOperation( diff --git a/packages/smithy-core/tests/unit/aio/test_client.py b/packages/smithy-core/tests/unit/aio/test_client.py index 8af452517..215d7ff14 100644 --- a/packages/smithy-core/tests/unit/aio/test_client.py +++ b/packages/smithy-core/tests/unit/aio/test_client.py @@ -1,7 +1,9 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 +from collections.abc import AsyncIterator, Callable from dataclasses import replace +from io import BytesIO, UnsupportedOperation from typing import Any, cast from unittest.mock import AsyncMock, Mock, patch @@ -9,7 +11,8 @@ from smithy_core.aio.client import AsyncClient from smithy_core.aio.eventstream import DuplexEventStream, InputEventStream from smithy_core.aio.interfaces import ClientTransport -from smithy_core.aio.retries import AdaptiveRetryStrategy +from smithy_core.aio.retries import AdaptiveRetryStrategy, SimpleRetryStrategy +from smithy_core.aio.types import AsyncBytesReader from smithy_core.auth import AuthOption from smithy_core.exceptions import CallError, SmithyError, UnsupportedTransportError from smithy_core.interceptors import ( @@ -19,10 +22,13 @@ OutputContext, ) from smithy_core.response import EMPTY_RESPONSE_METADATA +from smithy_core.retries import ExponentialRetryBackoffStrategy from smithy_core.shapes import ShapeID from smithy_core.types import TypedProperties from ._pipeline_harness import ( + DATA_STREAM_OPERATION, + OPERATION, DuplexTransport, NonDuplexTransport, StubEndpointResolver, @@ -30,6 +36,7 @@ StubEventReceiver, StubInput, StubOutput, + StubRequest, UndeclaredTransport, client_call, pipeline_harness, @@ -319,3 +326,116 @@ async def test_aexit_closes_client() -> None: assert client.closed is True assert client.transport.close_calls == 1 + + +class BodyRecordingTransport(NonDuplexTransport): + def __init__(self) -> None: + super().__init__() + self.bodies: list[bytes] = [] + + async def send(self, request: Any) -> Any: + self.bodies.append(await AsyncBytesReader(request.body).read()) + return await super().send(request) + + +async def _chunks() -> AsyncIterator[bytes]: + yield b"foobar" + + +def _streaming_call_failing_once( + harness: Any, body: Any, *, operation: Any = DATA_STREAM_OPERATION +) -> Any: + request = StubRequest() + request.body = body + harness.protocol.serialize_request = lambda **kwargs: request # type: ignore[method-assign] + + async def fail_first_attempt(**kwargs: Any) -> StubOutput: + if harness.protocol.deserialize_response_calls == 0: + harness.protocol.deserialize_response_calls += 1 + raise CallError("Service unavailable", is_retry_safe=True) + return StubOutput() + + harness.protocol.deserialize_response = fail_first_attempt # type: ignore[method-assign] + return replace( + client_call(), + operation=operation, + retry_strategy=SimpleRetryStrategy( + backoff_strategy=ExponentialRetryBackoffStrategy(backoff_scale_value=0), + max_attempts=2, + ), + ) + + +async def test_seekable_data_stream_is_retried_from_its_start() -> None: + harness = pipeline_harness(BodyRecordingTransport()) + body = BytesIO(b"foobar") + body.seek(2) + call = _streaming_call_failing_once(harness, AsyncBytesReader(body)) + + await harness.pipeline(call) + + # The caller's position is where the body starts, so retries don't resend 'fo'. + assert cast(BodyRecordingTransport, harness.transport).bodies == [b"obar", b"obar"] + + +class UnrewindableBody: + """A body that reports it's seekable but fails to seek.""" + + def seekable(self) -> bool: + return True + + def seek(self, offset: int, whence: int = 0) -> int: + raise UnsupportedOperation("seek") + + async def read(self, size: int = -1) -> bytes: + return b"foobar" + + +@pytest.mark.parametrize( + "factory", + [lambda: AsyncBytesReader(_chunks()), UnrewindableBody], + ids=["not-seekable", "fails-to-seek"], +) +async def test_data_stream_that_cant_be_rewound_is_sent_once( + factory: Callable[[], Any], +) -> None: + harness = pipeline_harness(BodyRecordingTransport()) + call = _streaming_call_failing_once(harness, factory()) + + with pytest.raises(CallError, match="Service unavailable"): + await harness.pipeline(call) + + assert cast(BodyRecordingTransport, harness.transport).bodies == [b"foobar"] + + +class SeekBackFailsBody(UnrewindableBody): + """A body that reports its position but fails to seek back to it.""" + + def seek(self, offset: int, whence: int = 0) -> int: + if whence == 1: + return 0 + raise UnsupportedOperation("seek") + + +async def test_data_stream_that_fails_to_rewind_returns_the_attempts_error() -> None: + harness = pipeline_harness(BodyRecordingTransport()) + call = _streaming_call_failing_once(harness, SeekBackFailsBody()) + + with pytest.raises(CallError, match="Service unavailable") as exc_info: + await harness.pipeline(call) + + assert cast(BodyRecordingTransport, harness.transport).bodies == [b"foobar"] + assert any( + note.startswith("Retry not attempted") for note in exc_info.value.__notes__ + ) + + +async def test_input_event_stream_is_not_retried() -> None: + harness = pipeline_harness(NonDuplexTransport()) + call = _streaming_call_failing_once(harness, BytesIO(b""), operation=OPERATION) + + stream = await harness.pipeline.input_stream(call, StubEvent) + + with pytest.raises(CallError, match="Service unavailable"): + await stream.await_output() + assert harness.transport.send_calls == 1 diff --git a/packages/smithy-http/.changes/next-release/smithy-http-bugfix-dbdfcc08b0ac4033a2f69774df0decdd.json b/packages/smithy-http/.changes/next-release/smithy-http-bugfix-dbdfcc08b0ac4033a2f69774df0decdd.json new file mode 100644 index 000000000..0a770cb71 --- /dev/null +++ b/packages/smithy-http/.changes/next-release/smithy-http-bugfix-dbdfcc08b0ac4033a2f69774df0decdd.json @@ -0,0 +1,4 @@ +{ + "type": "bugfix", + "description": "Fixed HTTP serializers sending seekable streaming payloads, such as files, without a Content-Length and sending a modeled Content-Length twice." +} diff --git a/packages/smithy-http/src/smithy_http/serializers.py b/packages/smithy-http/src/smithy_http/serializers.py index 471cb154b..182ae3562 100644 --- a/packages/smithy-http/src/smithy_http/serializers.py +++ b/packages/smithy-http/src/smithy_http/serializers.py @@ -1,11 +1,11 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 from base64 import b64encode -from collections.abc import Callable, Iterator, Sized +from collections.abc import Callable, Iterator from contextlib import AbstractContextManager, contextmanager from datetime import datetime from decimal import Decimal -from inspect import iscoroutinefunction +from inspect import iscoroutine from io import BytesIO from types import TracebackType from typing import TYPE_CHECKING @@ -34,7 +34,7 @@ from smithy_core.types import PathPattern, TimestampFormat from smithy_core.utils import serialize_float -from . import Field, tuples_to_fields +from . import Field, Fields, tuples_to_fields from .aio import HTTPRequest as _HTTPRequest from .aio import HTTPResponse as _HTTPResponse from .aio.interfaces import HTTPRequest, HTTPResponse @@ -110,26 +110,75 @@ def begin_struct(self, schema: Schema) -> Iterator[ShapeSerializer]: self.result = binding_serializer.build_request() -def _compute_content_length( - payload: object, -) -> int | None: - if (tell := getattr(payload, "tell", None)) is not None and not iscoroutinefunction( - tell - ): - start: int = tell() - if (end := _seek(payload, 0, 2)) is not None: - content_length: int = end - start - _seek(payload, start, 0) - return content_length - return None +def _set_content_length( + fields: Fields, + content_length: int | None, + raw_payload: "AsyncStreamingBlob | None", + *, + required: bool, +) -> None: + # A length the caller already set, such as with a modeled Content-Length + # member, takes precedence and isn't computed again. + if "content-length" in fields: + return + + # A body sent with Transfer-Encoding can't also have a Content-Length. + if "transfer-encoding" not in fields: + if raw_payload is not None: + content_length = _compute_content_length(raw_payload) + if content_length is not None: + fields.set_field(Field(name="content-length", values=[str(content_length)])) + return + if required: + raise SerializationError( + "This operation requires the content length of the input stream, but " + "it was not provided and was unable to be computed. Provide a " + "seekable stream or set the content length explicitly." + ) -def _seek(payload: object, pos: int, whence: int = 0) -> int | None: - if (seek := getattr(payload, "seek", None)) is not None and not iscoroutinefunction( - seek - ): - return seek(pos, whence) - return None + +def _compute_content_length(payload: "AsyncStreamingBlob") -> int | None: + """Determine the length of a raw payload without consuming it. + + ``bytes`` and ``bytearray`` report their length, and synchronously seekable + streams are measured from their current position to the end. The length of any + other stream, including async streams, is unknown. + """ + if isinstance(payload, bytes | bytearray): + return len(payload) + + tell = getattr(payload, "tell", None) + seek = getattr(payload, "seek", None) + if tell is None or seek is None: + return None + + try: + if (seekable := getattr(payload, "seekable", None)) is not None: + if _is_async(is_seekable := seekable()) or not is_seekable: + return None + if _is_async(start := tell()) or _is_async(seek(0, 2)): + return None + except (OSError, ValueError, NotImplementedError): + # Pipes raise io.UnsupportedOperation, closed files raise ValueError, and + # AsyncBytesReader raises NotImplementedError if its source can't tell. + return None + + # The stream is now at its end, so errors are raised instead of sending it on + # with an empty body. + end: int = tell() + seek(start, 0) + # A stream positioned past its end has nothing left to send. + return max(end - start, 0) + + +def _is_async(result: object) -> bool: + # Wrappers of async methods return a coroutine without being declared with + # ``async def``. Closing it avoids a "never awaited" warning. + if iscoroutine(result): + result.close() + return True + return False class HTTPBindingSerializer(InterceptingSerializer): @@ -171,7 +220,6 @@ def __init__( self._raw_payload_serializer: RawPayloadSerializer | None = None self._payload: AsyncBytesReader | AsyncBytesProvider | None = None self._content_type: str | None = payload_codec.media_type - self._content_length: int | None = None self._content_length_required = False self._writes_document_body = False self._result: HTTPRequest | None = None @@ -253,22 +301,22 @@ def build_request(self) -> HTTPRequest: self._body_context = None payload = self._payload + content_length: int | None = None + raw_payload: AsyncStreamingBlob | None = None if self._raw_payload_serializer is not None: raw_payload = self._raw_payload_serializer.payload - if isinstance(raw_payload, Sized): - self._content_length = len(raw_payload) - payload = AsyncBytesReader(raw_payload or b"") + if raw_payload is None and self._content_length_required: + # An unset payload is sent as an empty body, whose length is 0. + raw_payload = b"" elif self._sync_payload is not None: if ( self._binding_metadata.request.payload is not None or self._writes_document_body ): - self._content_length = self._sync_payload.tell() + content_length = self._sync_payload.tell() self._sync_payload.seek(0) payload = AsyncBytesReader(self._sync_payload) - assert payload is not None # noqa: S101 - headers = self.header_serializer.headers if self._content_type is not None and ( not self._binding_metadata.request.has_content_type_header @@ -276,18 +324,14 @@ def build_request(self) -> HTTPRequest: ): headers.append(("content-type", self._content_type)) - if self._content_length is not None: - headers.append(("content-length", str(self._content_length))) - fields = tuples_to_fields(headers) - if self._content_length_required and "content-length" not in fields: - content_length = _compute_content_length(payload) - if content_length is None: - raise SerializationError( - "This operation requires the content length of the input " - "stream, but it was not provided and was unable to be computed." - ) - fields.set_field(Field(name="content-length", values=[str(content_length)])) + _set_content_length( + fields, content_length, raw_payload, required=self._content_length_required + ) + if self._raw_payload_serializer is not None: + payload = AsyncBytesReader(raw_payload or b"") + + assert payload is not None # noqa: S101 self._result = _HTTPRequest( method=self._operation_metadata.method, @@ -388,8 +432,8 @@ def begin_struct(self, schema: Schema) -> Iterator[ShapeSerializer]: content_type: str | None = self._payload_codec.media_type content_length: int | None = None - content_length_source: object | None = None content_length_required = False + raw_payload: AsyncStreamingBlob | None = None binding_metadata = schema.get_extension(HTTP_BINDING_SCHEMA_EXTENSION) response_metadata = binding_metadata.response @@ -410,9 +454,8 @@ def begin_struct(self, schema: Schema) -> Iterator[ShapeSerializer]: ) yield binding_serializer raw_payload = payload_serializer.payload - if isinstance(raw_payload, Sized): - content_length = len(raw_payload) - content_length_source = raw_payload + if raw_payload is None and content_length_required: + raw_payload = b"" payload = AsyncBytesReader(raw_payload or b"") else: if payload_metadata.media_type is not None: @@ -454,20 +497,10 @@ def begin_struct(self, schema: Schema) -> Iterator[ShapeSerializer]: ): headers.append(("content-type", content_type)) - if content_length is not None: - headers.append(("content-length", str(content_length))) - fields = tuples_to_fields(headers) - if content_length_required and "content-length" not in fields: - content_length = _compute_content_length( - content_length_source if content_length_source is not None else payload - ) - if content_length is None: - raise SerializationError( - "This operation requires the content length of the input " - "stream, but it was not provided and was unable to be computed." - ) - fields.set_field(Field(name="content-length", values=[str(content_length)])) + _set_content_length( + fields, content_length, raw_payload, required=content_length_required + ) status = binding_serializer.response_code_serializer.response_code if status is None: diff --git a/packages/smithy-http/tests/unit/test_serializers.py b/packages/smithy-http/tests/unit/test_serializers.py index 9708ee2b7..c9550fe58 100644 --- a/packages/smithy-http/tests/unit/test_serializers.py +++ b/packages/smithy-http/tests/unit/test_serializers.py @@ -1,6 +1,7 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 import datetime +from collections.abc import AsyncIterator, Callable from dataclasses import dataclass, field from datetime import UTC from decimal import Decimal @@ -13,16 +14,18 @@ from smithy_core.aio.interfaces import StreamingBlob from smithy_core.aio.types import AsyncBytesReader from smithy_core.deserializers import DeserializeableShape, ShapeDeserializer +from smithy_core.exceptions import SerializationError from smithy_core.prelude import ( BIG_DECIMAL, BLOB, BOOLEAN, FLOAT, INTEGER, + LONG, STRING, TIMESTAMP, ) -from smithy_core.schemas import Schema +from smithy_core.schemas import MemberSchema, Schema from smithy_core.serializers import ( SerializeableStruct, ShapeSerializer, @@ -47,6 +50,7 @@ from smithy_http import Fields, tuples_to_fields from smithy_http.aio import HTTPResponse as _HTTPResponse from smithy_http.deserializers import HTTPResponseDeserializer +from smithy_http.interfaces import Fields as FieldsInterface from smithy_http.serializers import ( HTTPBindingSerializer, HTTPRequestSerializer, @@ -2309,3 +2313,222 @@ async def test_serialize_response_event_stream_creates_writeable_body() -> None: assert "header" in actual.fields assert actual.fields["header"].as_string() == "foo" + + +def _streaming_payload_schema(*, requires_length: bool = False) -> Schema: + payload: MemberSchema = { + "target": BLOB, + "traits": [HTTPPayloadTrait(), StreamingTrait()], + } + if requires_length: + payload["traits"].append(RequiresLengthTrait()) + return Schema.collection( + id=ShapeID( + "com.smithy#RequiredLengthStreamingPayload" + if requires_length + else "com.smithy#StreamingPayload" + ), + members={ + "payload": payload, + "contentLength": { + "target": LONG, + "traits": [HTTPHeaderTrait("Content-Length")], + }, + "transferEncoding": { + "target": STRING, + "traits": [HTTPHeaderTrait("Transfer-Encoding")], + }, + }, + ) + + +STREAMING_PAYLOAD = _streaming_payload_schema() +REQUIRED_LENGTH_STREAMING_PAYLOAD = _streaming_payload_schema(requires_length=True) +STREAMING_DATA = b"foobar" + + +class SeekableFalseBytesReader(NonSeekableBytesReader): + """A stream with seek and tell that reports it can't seek, like a pipe.""" + + def seekable(self) -> bool: + return False + + def seek(self, offset: int, whence: int = 0, /) -> int: + raise OSError("unseekable") + + def tell(self) -> int: + raise OSError("unseekable") + + +class SeekOnlyBytesReader(NonSeekableBytesReader): + """A stream that can seek but can't report its position.""" + + def seek(self, offset: int, whence: int = 0, /) -> int: + return self._wrapped.seek(offset, whence) + + +async def _async_chunks(data: bytes) -> AsyncIterator[bytes]: + yield data + + +def _serialize_streaming( + kind: str, + payload: StreamingBlob | None, + *, + content_length: int | None = None, + transfer_encoding: str | None = None, + schema: Schema = STREAMING_PAYLOAD, +) -> FieldsInterface: + """Serialize a request the way generated clients do, or a response, and return + its fields. A payload of None leaves the payload member unset.""" + + def write(serializer: ShapeSerializer) -> None: + if payload is not None: + serializer.write_data_stream(schema.members["payload"], payload) + if content_length is not None: + serializer.write_long(schema.members["contentLength"], content_length) + if transfer_encoding is not None: + serializer.write_string( + schema.members["transferEncoding"], transfer_encoding + ) + + http_trait = HTTPTrait({"method": "PUT", "code": 200, "uri": "/"}) + if kind == "request": + serializer = HTTPBindingSerializer( + payload_codec=JSONCodec(), schema=schema, http_trait=http_trait + ) + write(serializer) + return serializer.build_request().fields + + response_serializer = HTTPResponseSerializer( + payload_codec=JSONCodec(), http_trait=http_trait + ) + with response_serializer.begin_struct(schema) as struct_serializer: + write(struct_serializer) + assert response_serializer.result is not None + return response_serializer.result.fields + + +@pytest.mark.parametrize("kind", ["request", "response"]) +@pytest.mark.parametrize("offset", [0, 2]) +async def test_streaming_payload_sets_content_length(kind: str, offset: int) -> None: + payload = BytesIO(STREAMING_DATA) + payload.seek(offset) + fields = _serialize_streaming(kind, payload) + + # The length counts what's left from the current position, which is kept. + assert fields["content-length"].values == [str(len(STREAMING_DATA) - offset)] + assert payload.tell() == offset + + +@pytest.mark.parametrize("kind", ["request", "response"]) +async def test_streaming_payload_past_its_end_has_zero_content_length( + kind: str, +) -> None: + payload = BytesIO(STREAMING_DATA) + payload.seek(len(STREAMING_DATA) + 4) + fields = _serialize_streaming(kind, payload) + assert fields["content-length"].values == ["0"] + + +@pytest.mark.parametrize("kind", ["request", "response"]) +@pytest.mark.parametrize( + "factory", + [ + lambda: NonSeekableBytesReader(STREAMING_DATA), + lambda: SeekableFalseBytesReader(STREAMING_DATA), + lambda: AsyncBytesReader(BytesIO(STREAMING_DATA)), + lambda: _async_chunks(STREAMING_DATA), + lambda: AsyncBytesReader(SeekOnlyBytesReader(STREAMING_DATA)), + ], + ids=[ + "no-seek", + "seekable-false", + "async-seek", + "async-iterable", + "async-reader-without-tell", + ], +) +async def test_streaming_payload_of_unknown_length_omits_content_length( + kind: str, factory: Callable[[], StreamingBlob] +) -> None: + fields = _serialize_streaming(kind, factory()) + assert "content-length" not in fields + + +@pytest.mark.parametrize("kind", ["request", "response"]) +@pytest.mark.parametrize( + "schema", + [STREAMING_PAYLOAD, REQUIRED_LENGTH_STREAMING_PAYLOAD], + ids=["optional-length", "requires-length"], +) +@pytest.mark.parametrize( + "factory", + [lambda: STREAMING_DATA, lambda: NonSeekableBytesReader(STREAMING_DATA)], + ids=["bytes", "no-seek"], +) +async def test_modeled_content_length_is_used_once( + kind: str, schema: Schema, factory: Callable[[], StreamingBlob] +) -> None: + fields = _serialize_streaming(kind, factory(), content_length=6, schema=schema) + assert fields["content-length"].values == ["6"] + + +async def test_requires_length_computes_content_length() -> None: + fields = _serialize_streaming( + "request", BytesIO(STREAMING_DATA), schema=REQUIRED_LENGTH_STREAMING_PAYLOAD + ) + assert fields["content-length"].values == [str(len(STREAMING_DATA))] + + +@pytest.mark.parametrize("kind", ["request", "response"]) +@pytest.mark.parametrize( + "factory", + [ + lambda: NonSeekableBytesReader(STREAMING_DATA), + lambda: _async_chunks(STREAMING_DATA), + ], + ids=["no-seek", "async-iterable"], +) +async def test_requires_length_raises_for_stream_of_unknown_length( + kind: str, factory: Callable[[], StreamingBlob] +) -> None: + with pytest.raises(SerializationError): + _serialize_streaming(kind, factory(), schema=REQUIRED_LENGTH_STREAMING_PAYLOAD) + + +@pytest.mark.parametrize("kind", ["request", "response"]) +async def test_unset_streaming_payload_that_requires_length_has_zero_length( + kind: str, +) -> None: + fields = _serialize_streaming(kind, None, schema=REQUIRED_LENGTH_STREAMING_PAYLOAD) + assert fields["content-length"].values == ["0"] + + +@pytest.mark.parametrize("kind", ["request", "response"]) +async def test_unset_streaming_payload_omits_content_length(kind: str) -> None: + # Requests without content, such as GETs, shouldn't send a Content-Length. + fields = _serialize_streaming(kind, None) + assert "content-length" not in fields + + +@pytest.mark.parametrize("kind", ["request", "response"]) +async def test_streaming_payload_with_transfer_encoding_omits_content_length( + kind: str, +) -> None: + fields = _serialize_streaming( + kind, BytesIO(STREAMING_DATA), transfer_encoding="chunked" + ) + assert "content-length" not in fields + assert fields["transfer-encoding"].values == ["chunked"] + + +@pytest.mark.parametrize("kind", ["request", "response"]) +async def test_requires_length_with_transfer_encoding_raises(kind: str) -> None: + with pytest.raises(SerializationError): + _serialize_streaming( + kind, + BytesIO(STREAMING_DATA), + transfer_encoding="chunked", + schema=REQUIRED_LENGTH_STREAMING_PAYLOAD, + )