Add protocol support for AWS JSON-RPC

Makefile

@@ -14,10 +14,12 @@ build-java: ## Builds the Java code generation packages.
 	cd codegen && ./gradlew clean build
 
 
-test-protocols: ## Generates and runs the restJson1 protocol tests.
-	cd codegen && ./gradlew :protocol-test:build
-	uv pip install codegen/protocol-test/build/smithyprojections/protocol-test/rest-json-1/python-client-codegen
-	uv run pytest codegen/protocol-test/build/smithyprojections/protocol-test/rest-json-1/python-client-codegen
+test-protocols: ## Generates and runs protocol tests for all supported protocols.
+	cd codegen && ./gradlew :protocol-test:clean :protocol-test:build
+	@set -e; for projection_dir in codegen/protocol-test/build/smithyprojections/protocol-test/*/python-client-codegen; do \
+		uv pip install "$$projection_dir"; \
+		uv run pytest "$$projection_dir"; \
+	done
 
 
 lint-py: ## Runs linters and formatters on the python packages.

codegen/aws/core/build.gradle.kts

@@ -12,4 +12,5 @@ extra["moduleName"] = "software.amazon.smithy.python.aws.codegen"
 dependencies {
     implementation(project(":core"))
     implementation(libs.smithy.aws.traits)
+    implementation(libs.smithy.protocol.test.traits)
 }

codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/AwsJson10ProtocolGenerator.java

@@ -0,0 +1,76 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package software.amazon.smithy.python.aws.codegen;
+
+import java.util.Set;
+import software.amazon.smithy.aws.traits.protocols.AwsJson1_0Trait;
+import software.amazon.smithy.model.node.ArrayNode;
+import software.amazon.smithy.model.node.ObjectNode;
+import software.amazon.smithy.model.shapes.ShapeId;
+import software.amazon.smithy.python.codegen.ApplicationProtocol;
+import software.amazon.smithy.python.codegen.GenerationContext;
+import software.amazon.smithy.python.codegen.HttpProtocolTestGenerator;
+import software.amazon.smithy.python.codegen.SymbolProperties;
+import software.amazon.smithy.python.codegen.generators.ProtocolGenerator;
+import software.amazon.smithy.python.codegen.writer.PythonWriter;
+import software.amazon.smithy.utils.SmithyInternalApi;
+
+@SmithyInternalApi
+public final class AwsJson10ProtocolGenerator implements ProtocolGenerator {
+    private static final Set<String> TESTS_TO_SKIP = Set.of(
+            // TODO: support the request compression trait
+            // https://smithy.io/2.0/spec/behavior-traits.html#smithy-api-requestcompression-trait
+            "SDKAppliedContentEncoding_awsJson1_0",
+            "SDKAppendsGzipAndIgnoresHttpProvidedEncoding_awsJson1_0",
+
+            // TODO: Fix for both REST-JSON and JSON-RPC
+            "AwsJson10ClientPopulatesDefaultValuesInInput",
+            "AwsJson10ClientSkipsTopLevelDefaultValuesInInput",
+            "AwsJson10ClientUsesExplicitlyProvidedMemberValuesOverDefaults",
+            "AwsJson10ClientPopulatesDefaultsValuesWhenMissingInResponse",
+            "AwsJson10ClientIgnoresNonTopLevelDefaultsOnMembersWithClientOptional",
+            "AwsJson10ClientIgnoresDefaultValuesIfMemberValuesArePresentInResponse",
+
+            // TODO: support client error-correction behavior when the server
+            // omits required values in modeled error responses.
+            "AwsJson10ClientErrorCorrectsWhenServerFailsToSerializeRequiredValues",
+            "AwsJson10ClientErrorCorrectsWithDefaultValuesWhenServerFailsToSerializeRequiredValues");
+
+    @Override
+    public ShapeId getProtocol() {
+        return AwsJson1_0Trait.ID;
+    }
+
+    @Override
+    public ApplicationProtocol getApplicationProtocol(GenerationContext context) {
+        var service = context.settings().service(context.model());
+        var trait = service.expectTrait(AwsJson1_0Trait.class);
+        var config = ObjectNode.builder()
+                .withMember("http", ArrayNode.fromStrings(trait.getHttp()))
+                .withMember("eventStreamHttp", ArrayNode.fromStrings(trait.getEventStreamHttp()))
+                .build();
+        return ApplicationProtocol.createDefaultHttpApplicationProtocol(config);
+    }
+
+    @Override
+    public void initializeProtocol(GenerationContext context, PythonWriter writer) {
+        writer.addDependency(AwsPythonDependency.SMITHY_AWS_CORE.withOptionalDependencies("json"));
+        writer.addImport("smithy_aws_core.aio.protocols", "AwsJson10ClientProtocol");
+        var serviceSymbol = context.symbolProvider().toSymbol(context.settings().service(context.model()));
+        var serviceSchema = serviceSymbol.expectProperty(SymbolProperties.SCHEMA);
+        writer.write("AwsJson10ClientProtocol($T)", serviceSchema);
+    }
+
+    @Override
+    public void generateProtocolTests(GenerationContext context) {
+        context.writerDelegator().useFileWriter("./tests/test_protocol.py", "tests.test_protocol", writer -> {
+            new HttpProtocolTestGenerator(
+                    context,
+                    getProtocol(),
+                    writer,
+                    (shape, testCase) -> TESTS_TO_SKIP.contains(testCase.getId())).run();
+        });
+    }
+}

codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/AwsJson11ProtocolGenerator.java

@@ -0,0 +1,63 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package software.amazon.smithy.python.aws.codegen;
+
+import java.util.Set;
+import software.amazon.smithy.aws.traits.protocols.AwsJson1_1Trait;
+import software.amazon.smithy.model.node.ArrayNode;
+import software.amazon.smithy.model.node.ObjectNode;
+import software.amazon.smithy.model.shapes.ShapeId;
+import software.amazon.smithy.python.codegen.ApplicationProtocol;
+import software.amazon.smithy.python.codegen.GenerationContext;
+import software.amazon.smithy.python.codegen.HttpProtocolTestGenerator;
+import software.amazon.smithy.python.codegen.SymbolProperties;
+import software.amazon.smithy.python.codegen.generators.ProtocolGenerator;
+import software.amazon.smithy.python.codegen.writer.PythonWriter;
+import software.amazon.smithy.utils.SmithyInternalApi;
+
+@SmithyInternalApi
+public final class AwsJson11ProtocolGenerator implements ProtocolGenerator {
+    private static final Set<String> TESTS_TO_SKIP = Set.of(
+            // TODO: support the request compression trait
+            // https://smithy.io/2.0/spec/behavior-traits.html#smithy-api-requestcompression-trait
+            "SDKAppliedContentEncoding_awsJson1_1",
+            "SDKAppendsGzipAndIgnoresHttpProvidedEncoding_awsJson1_1");
+
+    @Override
+    public ShapeId getProtocol() {
+        return AwsJson1_1Trait.ID;
+    }
+
+    @Override
+    public ApplicationProtocol getApplicationProtocol(GenerationContext context) {
+        var service = context.settings().service(context.model());
+        var trait = service.expectTrait(AwsJson1_1Trait.class);
+        var config = ObjectNode.builder()
+                .withMember("http", ArrayNode.fromStrings(trait.getHttp()))
+                .withMember("eventStreamHttp", ArrayNode.fromStrings(trait.getEventStreamHttp()))
+                .build();
+        return ApplicationProtocol.createDefaultHttpApplicationProtocol(config);
+    }
+
+    @Override
+    public void initializeProtocol(GenerationContext context, PythonWriter writer) {
+        writer.addDependency(AwsPythonDependency.SMITHY_AWS_CORE.withOptionalDependencies("json"));
+        writer.addImport("smithy_aws_core.aio.protocols", "AwsJson11ClientProtocol");
+        var serviceSymbol = context.symbolProvider().toSymbol(context.settings().service(context.model()));
+        var serviceSchema = serviceSymbol.expectProperty(SymbolProperties.SCHEMA);
+        writer.write("AwsJson11ClientProtocol($T)", serviceSchema);
+    }
+
+    @Override
+    public void generateProtocolTests(GenerationContext context) {
+        context.writerDelegator().useFileWriter("./tests/test_protocol.py", "tests.test_protocol", writer -> {
+            new HttpProtocolTestGenerator(
+                    context,
+                    getProtocol(),
+                    writer,
+                    (shape, testCase) -> TESTS_TO_SKIP.contains(testCase.getId())).run();
+        });
+    }
+}

codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/AwsProtocolsIntegration.java

@@ -16,6 +16,6 @@
 public class AwsProtocolsIntegration implements PythonIntegration {
     @Override
     public List<ProtocolGenerator> getProtocolGenerators() {
-        return List.of();
+        return List.of(new AwsJson10ProtocolGenerator(), new AwsJson11ProtocolGenerator());
     }
 }

codegen/core/src/main/java/software/amazon/smithy/python/codegen/HttpProtocolTestGenerator.java

@@ -16,6 +16,7 @@
 import java.util.logging.Logger;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
+import software.amazon.smithy.aws.traits.auth.SigV4Trait;
 import software.amazon.smithy.codegen.core.CodegenException;
 import software.amazon.smithy.codegen.core.Symbol;
 import software.amazon.smithy.model.Model;
@@ -188,12 +189,14 @@ private void generateRequestTest(OperationShape operation, HttpRequestTestCase t
                                     endpoint_uri="https://$L/$L",
                                     transport = $T(),
                                     retry_strategy=SimpleRetryStrategy(max_attempts=1),
+                                    ${C|}
                                 )
                                 """,
                                 CodegenUtils.getConfigSymbol(context.settings()),
                                 host,
                                 path,
-                                REQUEST_TEST_ASYNC_HTTP_CLIENT_SYMBOL);
+                                REQUEST_TEST_ASYNC_HTTP_CLIENT_SYMBOL,
+                                (Runnable) this::writeSigV4TestConfig);
                     }));
 
                     // Generate the input using the expected shape and params
@@ -437,13 +440,15 @@ private void generateResponseTest(OperationShape operation, HttpResponseTestCase
                                         headers=$J,
                                         body=b$S,
                                     ),
+                                    ${C|}
                                 )
                                 """,
                                 CodegenUtils.getConfigSymbol(context.settings()),
                                 RESPONSE_TEST_ASYNC_HTTP_CLIENT_SYMBOL,
                                 testCase.getCode(),
                                 CodegenUtils.toTuples(testCase.getHeaders()),
-                                testCase.getBody().filter(body -> !body.isEmpty()).orElse(""));
+                                testCase.getBody().filter(body -> !body.isEmpty()).orElse(""),
+                                (Runnable) this::writeSigV4TestConfig);
                     }));
                     // Create an empty input object to pass
                     var inputShape = model.expectShape(operation.getInputShape(), StructureShape.class);
@@ -490,13 +495,15 @@ private void generateErrorResponseTest(
                                         headers=$J,
                                         body=b$S,
                                     ),
+                                    ${C|}
                                 )
                                 """,
                                 CodegenUtils.getConfigSymbol(context.settings()),
                                 RESPONSE_TEST_ASYNC_HTTP_CLIENT_SYMBOL,
                                 testCase.getCode(),
                                 CodegenUtils.toTuples(testCase.getHeaders()),
-                                testCase.getBody().orElse(""));
+                                testCase.getBody().orElse(""),
+                                (Runnable) this::writeSigV4TestConfig);
                     }));
                     // Create an empty input object to pass
                     var inputShape = model.expectShape(operation.getInputShape(), StructureShape.class);
@@ -531,7 +538,7 @@ private void assertResponseEqual(HttpMessageTestCase testCase, Shape operationOr
                 .findAny();
 
         if (streamBinding.isEmpty()) {
-            writer.write("assert actual == expected\n");
+            writer.write("_assert_modeled_value(actual, expected)\n");
             return;
         }
 
@@ -556,7 +563,7 @@ assert isinstance(actual.$1L, AsyncByteStream)
                 compareMediaBlob(testCase, writer);
                 continue;
             }
-            writer.write("assert actual.$1L == expected.$1L\n", memberName);
+            writer.write("_assert_modeled_value(actual.$1L, expected.$1L)\n", memberName);
         }
     }
 
@@ -607,10 +614,26 @@ private void writeClientBlock(
         });
     }
 
+    private void writeSigV4TestConfig() {
+        if (!service.hasTrait(SigV4Trait.class)) {
+            return;
+        }
+        writer.addImport("smithy_aws_core.identity", "StaticCredentialsResolver");
+        writer.write("""
+                region="us-east-1",
+                aws_access_key_id="test-access-key-id",
+                aws_secret_access_key="test-secret-access-key",
+                aws_credentials_identity_resolver=StaticCredentialsResolver(),
+                """);
+    }
+
     private void writeUtilStubs(Symbol serviceSymbol) {
         LOGGER.fine(String.format("Writing utility stubs for %s : %s", serviceSymbol.getName(), protocol.getName()));
         writer.addDependency(SmithyPythonDependency.SMITHY_CORE);
         writer.addDependency(SmithyPythonDependency.SMITHY_HTTP);
+        writer.addStdlibImport("dataclasses", "fields");
+        writer.addStdlibImport("dataclasses", "is_dataclass");
+        writer.addStdlibImport("math", "isnan");
         writer.addImports("smithy_http.interfaces",
                 Set.of(
                         "HTTPRequestConfiguration",
@@ -621,6 +644,24 @@ private void writeUtilStubs(Symbol serviceSymbol) {
         writer.addImport("smithy_core.aio.utils", "async_list");
 
         writer.write("""
+                def _assert_modeled_value(actual: object, expected: object) -> None:
+                    if isinstance(expected, float) and isnan(expected):
+                        assert isinstance(actual, float)
+                        assert isnan(actual)
+                        return
+
+                    if is_dataclass(expected):
+                        assert is_dataclass(actual)
+                        for field in fields(expected):
+                            _assert_modeled_value(
+                                getattr(actual, field.name),
+                                getattr(expected, field.name),
+                            )
+                        return
+
+                    assert actual == expected
+
+
                 class $1L($2T):
                     ""\"A test error that subclasses the service-error for protocol tests.""\"
 

codegen/core/src/main/java/software/amazon/smithy/python/codegen/integrations/RestJsonProtocolGenerator.java

@@ -31,12 +31,6 @@
 public class RestJsonProtocolGenerator implements ProtocolGenerator {
 
     private static final Set<String> TESTS_TO_SKIP = Set.of(
-            // These two tests essentially try to assert nan == nan,
-            // which is never true. We should update the generator to
-            // make specific assertions for these.
-            "RestJsonSupportsNaNFloatHeaderOutputs",
-            "RestJsonSupportsNaNFloatInputs",
-
             // This requires support of idempotency autofill
             "RestJsonQueryIdempotencyTokenAutoFill",
 

codegen/protocol-test/build.gradle.kts

@@ -30,5 +30,6 @@ repositories {
 
 dependencies {
     implementation(project(":core"))
+    implementation(project(":aws:core"))
     implementation(libs.smithy.aws.protocol.tests)
 }

codegen/protocol-test/smithy-build.json

@@ -22,6 +22,50 @@
                     "moduleVersion": "0.0.1"
                 }
             }
+        },
+        "aws-json-1-0": {
+            "transforms": [
+                {
+                    "name": "includeServices",
+                    "args": {
+                        "services": [
+                            "aws.protocoltests.json10#JsonRpc10"
+                        ]
+                    }
+                },
+                {
+                    "name": "removeUnusedShapes"
+                }
+            ],
+            "plugins": {
+                "python-client-codegen": {
+                    "service": "aws.protocoltests.json10#JsonRpc10",
+                    "module": "awsjson10",
+                    "moduleVersion": "0.0.1"
+                }
+            }
+        },
+        "aws-json-1-1": {
+            "transforms": [
+                {
+                    "name": "includeServices",
+                    "args": {
+                        "services": [
+                            "aws.protocoltests.json#JsonProtocol"
+                        ]
+                    }
+                },
+                {
+                    "name": "removeUnusedShapes"
+                }
+            ],
+            "plugins": {
+                "python-client-codegen": {
+                    "service": "aws.protocoltests.json#JsonProtocol",
+                    "module": "awsjson11",
+                    "moduleVersion": "0.0.1"
+                }
+            }
         }
     }
 }