From 8a8aa01f8e4d374986b7323e081006b9c1e880e6 Mon Sep 17 00:00:00 2001
From: Aaron J Todd <todaaron@amazon.com>
Date: Mon, 25 Sep 2023 09:37:28 -0400
Subject: [PATCH] feat(codegen): add protocol tests for client side error
 correction

---
 codegen/protocol-tests/build.gradle.kts       |   4 +
 .../model/error-correction-tests.smithy       | 156 ++++++++++++++++++
 2 files changed, 160 insertions(+)
 create mode 100644 codegen/protocol-tests/model/error-correction-tests.smithy

diff --git a/codegen/protocol-tests/build.gradle.kts b/codegen/protocol-tests/build.gradle.kts
index a3e74ecc2daa..42d14f7968c2 100644
--- a/codegen/protocol-tests/build.gradle.kts
+++ b/codegen/protocol-tests/build.gradle.kts
@@ -35,6 +35,10 @@ val enabledProtocols = listOf(
     ProtocolTest("apigateway", "com.amazonaws.apigateway#BackplaneControlService"),
     ProtocolTest("glacier", "com.amazonaws.glacier#Glacier"),
     ProtocolTest("machinelearning", "com.amazonaws.machinelearning#AmazonML_20141212", sdkId = "Machine Learning"),
+
+    // Custom hand written tests
+    ProtocolTest("error-correction-json", "aws.protocoltests.errorcorrection#RequiredValueJson"),
+    ProtocolTest("error-correction-xml", "aws.protocoltests.errorcorrection#RequiredValueXml"),
 )
 
 codegen {
diff --git a/codegen/protocol-tests/model/error-correction-tests.smithy b/codegen/protocol-tests/model/error-correction-tests.smithy
new file mode 100644
index 000000000000..13ad29f921a4
--- /dev/null
+++ b/codegen/protocol-tests/model/error-correction-tests.smithy
@@ -0,0 +1,156 @@
+$version: "2.0"
+
+namespace aws.protocoltests.errorcorrection
+
+use aws.api#service
+use aws.protocols#awsJson1_0
+use aws.protocols#restXml
+use smithy.test#httpResponseTests
+
+@service(sdkId: "Error Correction Json")
+@awsJson1_0
+service RequiredValueJson {
+    operations: [SayHello],
+    version: "1"
+}
+
+
+@service(sdkId: "Error Correction Xml")
+@restXml
+service RequiredValueXml {
+    operations: [SayHelloXml],
+    version: "1"
+}
+
+@error("client")
+structure Error {
+    @required
+    requestId: String
+
+    @required
+    message: String
+}
+
+@http(method: "POST", uri: "/")
+operation SayHello { output: TestOutputDocument, errors: [Error] }
+
+@http(method: "POST", uri: "/")
+operation SayHelloXml { output: TestOutput, errors: [Error] }
+
+structure TestOutputDocument with [TestStruct] { innerField: Nested, @required document: Document }
+structure TestOutput with [TestStruct] { innerField: Nested }
+
+@mixin
+structure TestStruct {
+    @required
+    foo: String,
+
+    @required
+    byteValue: Byte,
+
+    @required
+    intValue: Integer,
+
+    @required
+    listValue: StringList,
+
+    @required
+    mapValue: ListMap,
+
+    @required
+    nestedListValue: NestedList
+
+    @required
+    nested: Nested
+
+    @required
+    blob: Blob
+
+    @required
+    enum: MyEnum
+
+    @required
+    union: MyUnion
+
+    notRequired: String
+
+    @required
+    timestampValue: Timestamp
+}
+
+enum MyEnum {
+    A,
+    B,
+    C
+}
+
+union MyUnion {
+    A: Integer,
+    B: String,
+    C: Unit
+}
+
+structure Nested {
+    @required
+    a: String
+}
+
+list StringList {
+    member: String
+}
+
+list NestedList {
+    member: StringList
+}
+
+map ListMap {
+    key: String,
+    value: StringList
+}
+
+// NOTE: there is no way to model enum or union defaults in an `httpResponseTest` because the default is the generated
+// "SdkUnknown" variant.
+apply SayHello @httpResponseTests([
+    {
+        id: "error_recovery_json",
+        protocol: awsJson1_0,
+        params: {
+            union: { A: 5 },
+            enum: "A",
+            foo: "",
+            byteValue: 0,
+            intValue: 0,
+            blob: "",
+            listValue: [],
+            mapValue: {},
+            nestedListValue: [],
+            document: null,
+            nested: { a: "" },
+            timestampValue: 0
+        },
+        code: 200,
+        body: "{\"union\": { \"A\": 5 }, \"enum\": \"A\" }"
+    }
+])
+
+apply SayHelloXml @httpResponseTests([
+    {
+        id: "error_recovery_xml",
+        protocol: restXml,
+        params: {
+            union: { A: 5 },
+            enum: "A",
+            foo: "",
+            byteValue: 0,
+            intValue: 0,
+            blob: "",
+            listValue: [],
+            mapValue: {},
+            nestedListValue: [],
+            nested: { a: "" },
+            timestampValue: 0
+        },
+        code: 200,
+        body: "<TestOutput><union><A>5</A></union><enum>A</enum></TestOutput>"
+    }
+])
\ No newline at end of file