comments

Author: haydenbaker

packages/smithy-aws-core/src/smithy_aws_core/aio/__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/src/smithy_aws_core/protocols/restjson.py renamed to packages/smithy-aws-core/src/smithy_aws_core/aio/protocols.py

@@ -1,7 +1,7 @@
 from typing import Literal
 
 from smithy_aws_core.traits import RestJson1Trait
-from smithy_core.protocols import HttpBindingClientProtocol
+from smithy_http.aio.protocols import HttpBindingClientProtocol
 from smithy_core.codecs import Codec
 from smithy_core.shapes import ShapeID
 from smithy_json import JSONCodec

packages/smithy-core/src/smithy_core/protocols.py renamed to packages/smithy-http/src/smithy_http/aio/protocols.py

@@ -29,6 +29,8 @@
     TimestampFormatTrait,
     EndpointTrait,
     HTTPErrorTrait,
+    MediaTypeTrait,
+    StreamingTrait,
 )
 from smithy_core.shapes import ShapeType
 from smithy_core.utils import serialize_float
@@ -83,22 +85,33 @@ def begin_struct(self, schema: Schema) -> Iterator[ShapeSerializer]:
         if self._endpoint_trait is not None:
             host_prefix = self._endpoint_trait.host_prefix
 
+        content_type = self._payload_codec.media_type
+
         if (payload_member := self._get_payload_member(schema)) is not None:
             if payload_member.shape_type in (ShapeType.BLOB, ShapeType.STRING):
+                content_type = (
+                    "application/octet-stream"
+                    if payload_member.shape_type is ShapeType.BLOB
+                    else "text/plain"
+                )
                 payload_serializer = RawPayloadSerializer()
                 binding_serializer = HTTPRequestBindingSerializer(
                     payload_serializer, self._http_trait.path, host_prefix
                 )
                 yield binding_serializer
                 payload = payload_serializer.payload
             else:
+                if (media_type := payload_member.get_trait(MediaTypeTrait)) is not None:
+                    content_type = media_type.value
                 payload = BytesIO()
                 payload_serializer = self._payload_codec.create_serializer(payload)
                 binding_serializer = HTTPRequestBindingSerializer(
                     payload_serializer, self._http_trait.path, host_prefix
                 )
                 yield binding_serializer
         else:
+            if self._get_eventstreaming_member(schema) is not None:
+                content_type = "application/vnd.amazon.eventstream"
             payload = BytesIO()
             payload_serializer = self._payload_codec.create_serializer(payload)
             with payload_serializer.begin_struct(schema) as body_serializer:
@@ -112,9 +125,9 @@ def begin_struct(self, schema: Schema) -> Iterator[ShapeSerializer]:
         ) is not None and not iscoroutinefunction(seek):
             seek(0)
 
-        # TODO: conditional on empty-ness and a param of the protocol?
+        # TODO: conditional on empty-ness and based on the protocol
         headers = binding_serializer.header_serializer.headers
-        headers.append(("content-type", self._payload_codec.media_type))
+        headers.append(("content-type", content_type))
 
         self.result = _HTTPRequest(
             method=self._http_trait.method,
@@ -136,6 +149,15 @@ def _get_payload_member(self, schema: Schema) -> Schema | None:
                 return member
         return None
 
+    def _get_eventstreaming_member(self, schema: Schema) -> Schema | None:
+        for member in schema.members.values():
+            if (
+                member.get_trait(StreamingTrait) is not None
+                and member.shape_type is ShapeType.UNION
+            ):
+                return member
+        return None
+
 
 class HTTPRequestBindingSerializer(InterceptingSerializer):
     """Delegates HTTP request bindings to binding-location-specific serializers."""

packages/smithy-http/src/smithy_http/serializers.py

@@ -1572,11 +1572,16 @@ def payload_cases() -> list[HTTPMessageTestCase]:
         ),
         HTTPMessageTestCase(
             HTTPStringPayload(payload="foo"),
-            HTTPMessage(body=b"foo"),
+            HTTPMessage(
+                fields=tuples_to_fields([("content-type", "text/plain")]), body=b"foo"
+            ),
         ),
         HTTPMessageTestCase(
             HTTPBlobPayload(payload=b"\xde\xad\xbe\xef"),
-            HTTPMessage(body=b"\xde\xad\xbe\xef"),
+            HTTPMessage(
+                fields=tuples_to_fields([("content-type", "application/octet-stream")]),
+                body=b"\xde\xad\xbe\xef",
+            ),
         ),
         HTTPMessageTestCase(
             HTTPStructuredPayload(payload=HTTPStringPayload(payload="foo")),
@@ -1589,7 +1594,10 @@ def async_streaming_payload_cases() -> list[HTTPMessageTestCase]:
     return [
         HTTPMessageTestCase(
             HTTPStreamingPayload(payload=AsyncBytesReader(b"\xde\xad\xbe\xef")),
-            HTTPMessage(body=AsyncBytesReader(b"\xde\xad\xbe\xef")),
+            HTTPMessage(
+                fields=tuples_to_fields([("content-type", "application/octet-stream")]),
+                body=AsyncBytesReader(b"\xde\xad\xbe\xef"),
+            ),
         ),
     ]
 
@@ -1625,8 +1633,10 @@ async def test_serialize_http_request(case: HTTPMessageTestCase) -> None:
     actual_query = actual.destination.query or ""
     expected_query = case.request.destination.query or ""
     assert actual_query == expected_query
-    # set the content-type field here, otherwise cases would have to duplicate it everywhere
-    expected.fields.set_field(CONTENT_TYPE_FIELD)
+    # set the content-type field here, otherwise cases would have to duplicate it everywhere,
+    # but if the field is already set in the case, don't override it
+    if expected.fields.get(CONTENT_TYPE_FIELD.name) is None:
+        expected.fields.set_field(CONTENT_TYPE_FIELD)
     assert actual.fields == expected.fields
 
     if case.request.body:
@@ -1651,6 +1661,9 @@ async def test_serialize_http_response(case: HTTPMessageTestCase) -> None:
     expected = case.request
 
     assert actual is not None
+    # Remove content-type from expected, we're re-using the request cases for brevity
+    if expected.fields.get(CONTENT_TYPE_FIELD.name) is not None:
+        del expected.fields[CONTENT_TYPE_FIELD.name]
     assert actual.fields == expected.fields
     assert actual.status == expected.status