Test on all JSON supported protocols

Author: Fahad Zubair

codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/testutil/CodegenIntegrationTest.kt

@@ -5,10 +5,19 @@
 
 package software.amazon.smithy.rust.codegen.core.testutil
 
+import software.amazon.smithy.aws.traits.protocols.AwsJson1_0Trait
+import software.amazon.smithy.aws.traits.protocols.AwsJson1_1Trait
+import software.amazon.smithy.aws.traits.protocols.RestJson1Trait
+import software.amazon.smithy.aws.traits.protocols.RestXmlTrait
 import software.amazon.smithy.build.PluginContext
 import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.node.ObjectNode
 import software.amazon.smithy.model.node.ToNode
+import software.amazon.smithy.model.shapes.ServiceShape
+import software.amazon.smithy.model.shapes.ShapeId
+import software.amazon.smithy.model.traits.AbstractTrait
+import software.amazon.smithy.model.transform.ModelTransformer
+import software.amazon.smithy.protocol.traits.Rpcv2CborTrait
 import software.amazon.smithy.rust.codegen.core.smithy.RuntimeConfig
 import software.amazon.smithy.rust.codegen.core.util.runCommand
 import java.io.File
@@ -153,3 +162,128 @@ fun codegenIntegrationTest(
     logger.fine(out.toString())
     return testDir
 }
+
+/**
+ * Metadata associated with a protocol that provides additional information needed for testing.
+ *
+ * @property protocol The protocol enum value this metadata is associated with
+ * @property contentType The HTTP Content-Type header value associated with this protocol.
+ */
+data class ProtocolMetadata(
+    val protocol: ModelProtocol,
+    val contentType: String,
+)
+
+/**
+ * Represents the supported protocol traits in Smithy models.
+ *
+ * @property trait The Smithy trait instance with which the service shape must be annotated.
+ */
+enum class ModelProtocol(val trait: AbstractTrait) {
+    AwsJson10(AwsJson1_0Trait.builder().build()),
+    AwsJson11(AwsJson1_1Trait.builder().build()),
+    RestJson(RestJson1Trait.builder().build()),
+    RestXml(RestXmlTrait.builder().build()),
+    Rpcv2Cbor(Rpcv2CborTrait.builder().build()),
+    ;
+
+    // Create metadata after enum is initialized
+    val metadata: ProtocolMetadata by lazy {
+        when (this) {
+            AwsJson10 -> ProtocolMetadata(this, "application/x-amz-json-1.0")
+            AwsJson11 -> ProtocolMetadata(this, "application/x-amz-json-1.1")
+            RestJson -> ProtocolMetadata(this, "application/json")
+            RestXml -> ProtocolMetadata(this, "application/xml")
+            Rpcv2Cbor -> ProtocolMetadata(this, "application/cbor")
+        }
+    }
+
+    companion object {
+        private val TRAIT_IDS = values().map { it.trait.toShapeId() }.toSet()
+        val ALL: Set<ModelProtocol> = values().toSet()
+
+        fun getTraitIds() = TRAIT_IDS
+    }
+}
+
+/**
+ * Removes all existing protocol traits annotated on the given service,
+ * then sets the provided `protocol` as the sole protocol trait for the service.
+ */
+fun Model.replaceProtocolTraitOnServerShapeId(
+    serviceShapeId: ShapeId,
+    modelProtocol: ModelProtocol,
+): Model {
+    val serviceShape = this.expectShape(serviceShapeId, ServiceShape::class.java)
+    return replaceProtocolTraitOnServiceShape(serviceShape, modelProtocol)
+}
+
+/**
+ * Removes all existing protocol traits annotated on the given service shape,
+ * then sets the provided `protocol` as the sole protocol trait for the service.
+ */
+fun Model.replaceProtocolTraitOnServiceShape(
+    serviceShape: ServiceShape,
+    modelProtocol: ModelProtocol,
+): Model {
+    val serviceBuilder = serviceShape.toBuilder()
+    ModelProtocol.getTraitIds().forEach { traitId ->
+        serviceBuilder.removeTrait(traitId)
+    }
+    val service = serviceBuilder.addTrait(modelProtocol.trait).build()
+    return ModelTransformer.create().replaceShapes(this, listOf(service))
+}
+
+/**
+ * Processes a Smithy model string by applying different protocol traits and invoking the tests block on the model.
+ * For each protocol, this function:
+ *  1. Parses the Smithy model string
+ *  2. Replaces any existing protocol traits on service shapes with the specified protocol
+ *  3. Runs the provided test with the transformed model and protocol metadata
+ *
+ * @param protocolTraitIds Set of protocols to test against
+ * @param test Function that receives the transformed model and protocol metadata for testing
+ */
+fun String.forProtocols(
+    protocolTraitIds: Set<ModelProtocol>,
+    test: (Model, ProtocolMetadata) -> Unit,
+) {
+    val baseModel = this.asSmithyModel(smithyVersion = "2")
+    val serviceShapes = baseModel.serviceShapes.toList()
+
+    protocolTraitIds.forEach { protocol ->
+        val transformedModel =
+            serviceShapes.fold(baseModel) { acc, shape ->
+                acc.replaceProtocolTraitOnServiceShape(shape, protocol)
+            }
+        test(transformedModel, protocol.metadata)
+    }
+}
+
+/**
+ * Convenience overload that accepts vararg protocols instead of a Set.
+ *
+ * @param protocols Variable number of protocols to test against
+ * @param test Function that receives the transformed model and protocol metadata for testing
+ * @see forProtocols
+ */
+fun String.forProtocols(
+    vararg protocols: ModelProtocol,
+    test: (Model, ProtocolMetadata) -> Unit,
+) {
+    forProtocols(protocols.toSet(), test)
+}
+
+/**
+ * Tests a Smithy model string against all supported protocols, with optional exclusions.
+ *
+ * @param exclude Set of protocols to exclude from testing (default is empty)
+ * @param test Function that receives the transformed model and protocol metadata for testing
+ * @see forProtocols
+ */
+fun String.forAllProtocols(
+    exclude: Set<ModelProtocol> = emptySet(),
+    test: (Model, ProtocolMetadata) -> Unit,
+) {
+    forProtocols(ModelProtocol.ALL - exclude, test)
+}

codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/protocol/ServerProtocol.kt

@@ -134,7 +134,7 @@ fun jsonParserGenerator(
     codegenContext: ServerCodegenContext,
     httpBindingResolver: HttpBindingResolver,
     jsonName: (MemberShape) -> String,
-    additionalParserCustomizations: List<JsonParserCustomization> = listOf()
+    additionalParserCustomizations: List<JsonParserCustomization> = listOf(),
 ): JsonParserGenerator =
     JsonParserGenerator(
         codegenContext,
@@ -152,7 +152,6 @@ fun jsonParserGenerator(
             } else {
                 RuntimeType.smithyJson(codegenContext.runtimeConfig)
             },
-        ,
     )
 
 class ServerAwsJsonProtocol(

codegen-server/src/test/kotlin/software/amazon/smithy/rust/codegen/server/smithy/ConstraintsTest.kt

@@ -9,10 +9,6 @@ import io.kotest.inspectors.forAll
 import io.kotest.matchers.ints.shouldBeGreaterThan
 import io.kotest.matchers.shouldBe
 import org.junit.jupiter.api.Test
-import software.amazon.smithy.aws.traits.protocols.AwsJson1_0Trait
-import software.amazon.smithy.aws.traits.protocols.AwsJson1_1Trait
-import software.amazon.smithy.aws.traits.protocols.RestJson1Trait
-import software.amazon.smithy.aws.traits.protocols.RestXmlTrait
 import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.shapes.BooleanShape
 import software.amazon.smithy.model.shapes.ListShape
@@ -22,35 +18,27 @@ import software.amazon.smithy.model.shapes.ServiceShape
 import software.amazon.smithy.model.shapes.ShapeId
 import software.amazon.smithy.model.shapes.StringShape
 import software.amazon.smithy.model.shapes.StructureShape
-import software.amazon.smithy.model.traits.AbstractTrait
 import software.amazon.smithy.model.transform.ModelTransformer
-import software.amazon.smithy.protocol.traits.Rpcv2CborTrait
 import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
 import software.amazon.smithy.rust.codegen.core.smithy.RustCrate
 import software.amazon.smithy.rust.codegen.core.testutil.IntegrationTestParams
+import software.amazon.smithy.rust.codegen.core.testutil.ModelProtocol
 import software.amazon.smithy.rust.codegen.core.testutil.ServerAdditionalSettings
 import software.amazon.smithy.rust.codegen.core.testutil.asSmithyModel
+import software.amazon.smithy.rust.codegen.core.testutil.replaceProtocolTraitOnServerShapeId
 import software.amazon.smithy.rust.codegen.core.testutil.unitTest
 import software.amazon.smithy.rust.codegen.core.util.lookup
 import software.amazon.smithy.rust.codegen.server.smithy.testutil.serverIntegrationTest
 import software.amazon.smithy.rust.codegen.server.smithy.testutil.serverTestSymbolProvider
 import java.io.File
 
-enum class ModelProtocol(val trait: AbstractTrait) {
-    AwsJson10(AwsJson1_0Trait.builder().build()),
-    AwsJson11(AwsJson1_1Trait.builder().build()),
-    RestJson(RestJson1Trait.builder().build()),
-    RestXml(RestXmlTrait.builder().build()),
-    Rpcv2Cbor(Rpcv2CborTrait.builder().build()),
-}
-
 /**
  * Returns the Smithy constraints model from the common repository, with the specified protocol
  * applied to the service.
  */
 fun loadSmithyConstraintsModelForProtocol(modelProtocol: ModelProtocol): Pair<Model, ShapeId> {
     val (model, serviceShapeId) = loadSmithyConstraintsModel()
-    return Pair(model.replaceProtocolTrait(serviceShapeId, modelProtocol), serviceShapeId)
+    return Pair(model.replaceProtocolTraitOnServerShapeId(serviceShapeId, modelProtocol), serviceShapeId)
 }
 
 /**
@@ -65,23 +53,6 @@ fun loadSmithyConstraintsModel(): Pair<Model, ShapeId> {
     return Pair(model, serviceShapeId)
 }
 
-/**
- * Removes all existing protocol traits annotated on the given service,
- * then sets the provided `protocol` as the sole protocol trait for the service.
- */
-fun Model.replaceProtocolTrait(
-    serviceShapeId: ShapeId,
-    modelProtocol: ModelProtocol,
-): Model {
-    val serviceBuilder =
-        this.expectShape(serviceShapeId, ServiceShape::class.java).toBuilder()
-    for (p in ModelProtocol.values()) {
-        serviceBuilder.removeTrait(p.trait.toShapeId())
-    }
-    val service = serviceBuilder.addTrait(modelProtocol.trait).build()
-    return ModelTransformer.create().replaceShapes(this, listOf(service))
-}
-
 fun List<ShapeId>.containsAnyShapeId(ids: Collection<ShapeId>): Boolean {
     return ids.any { id -> this.any { shape -> shape == id } }
 }