Skip to content

Commit

Permalink
Test on all JSON supported protocols
Browse files Browse the repository at this point in the history
  • Loading branch information
Fahad Zubair committed Feb 4, 2025
1 parent 5712eed commit 726d808
Show file tree
Hide file tree
Showing 7 changed files with 211 additions and 99 deletions.
2 changes: 1 addition & 1 deletion .changelog/1738663351.md
Original file line number Diff line number Diff line change
Expand Up @@ -8,4 +8,4 @@ breaking: false
new_feature: true
bug_fix: false
---
Allow invalid UTF-8 characters to be replaced with
Enhanced UTF-8 handling: When `replaceInvalidUtf8` codegen flag is enabled, invalid UTF-8 sequences are now automatically replaced with the Unicode replacement character (U+FFFD) instead of causing errors
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
}
Original file line number Diff line number Diff line change
Expand Up @@ -153,7 +153,6 @@ fun jsonParserGenerator(
} else {
RuntimeType.smithyJson(codegenContext.runtimeConfig)
},
,
)

class ServerAwsJsonProtocol(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
}

/**
Expand All @@ -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 } }
}
Expand Down
Loading

0 comments on commit 726d808

Please sign in to comment.