All query tests passing!

alextwoods · alextwoods · commit be08f95d8866 · 2025-07-08T09:12:00.000-07:00
diff --git a/test/protocol-tests-core/src/main/java/software/amazon/awssdk/protocol/asserts/marshalling/CompositeMarshallingAssertion.java b/test/protocol-tests-core/src/main/java/software/amazon/awssdk/protocol/asserts/marshalling/CompositeMarshallingAssertion.java
@@ -18,7 +18,6 @@
 import com.github.tomakehurst.wiremock.verification.LoggedRequest;
 import java.util.ArrayList;
 import java.util.List;
-import software.amazon.awssdk.http.SdkHttpRequest;
 
 /**
  * Composite for MarshallingAssertion objects.
diff --git a/test/protocol-tests-core/src/main/java/software/amazon/awssdk/protocol/asserts/marshalling/MarshallingAssertion.java b/test/protocol-tests-core/src/main/java/software/amazon/awssdk/protocol/asserts/marshalling/MarshallingAssertion.java
@@ -16,7 +16,6 @@
 package software.amazon.awssdk.protocol.asserts.marshalling;
 
 import com.github.tomakehurst.wiremock.verification.LoggedRequest;
-import software.amazon.awssdk.http.SdkHttpRequest;
 
 /**
  * Assertion on the marshalled request.
diff --git a/test/protocol-tests-core/src/main/java/software/amazon/awssdk/protocol/runners/MarshallingTestRunner.java b/test/protocol-tests-core/src/main/java/software/amazon/awssdk/protocol/runners/MarshallingTestRunner.java
@@ -15,24 +15,27 @@
 
 package software.amazon.awssdk.protocol.runners;
 
+
+import static com.github.tomakehurst.wiremock.client.WireMock.aResponse;
+import static com.github.tomakehurst.wiremock.client.WireMock.any;
+import static com.github.tomakehurst.wiremock.client.WireMock.stubFor;
+import static com.github.tomakehurst.wiremock.client.WireMock.urlMatching;
 import static org.junit.Assert.assertEquals;
 
 import com.fasterxml.jackson.databind.JsonNode;
+import com.github.tomakehurst.wiremock.client.ResponseDefinitionBuilder;
+import com.github.tomakehurst.wiremock.client.WireMock;
 import com.github.tomakehurst.wiremock.verification.LoggedRequest;
 import java.lang.reflect.InvocationTargetException;
 import java.net.URI;
 import java.util.List;
 import org.junit.Assert;
 import software.amazon.awssdk.awscore.AwsRequest;
-import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
 import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
-import software.amazon.awssdk.core.SdkPlugin;
-import software.amazon.awssdk.core.SdkServiceClientConfiguration;
 import software.amazon.awssdk.core.interceptor.Context;
 import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
 import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
 import software.amazon.awssdk.core.sync.RequestBody;
-import software.amazon.awssdk.http.SdkHttpFullRequest;
 import software.amazon.awssdk.http.SdkHttpRequest;
 import software.amazon.awssdk.protocol.model.TestCase;
 import software.amazon.awssdk.protocol.reflect.ClientReflector;
@@ -47,52 +50,50 @@ class MarshallingTestRunner {
 
     private final IntermediateModel model;
     private final ClientReflector clientReflector;
-    private final RequestRecordingInterceptor recordingInterceptor;
+    private final LocalhostOnlyForWiremockInterceptor localhostOnlyForWiremockInterceptor;
 
     MarshallingTestRunner(IntermediateModel model, ClientReflector clientReflector) {
         this.model = model;
         this.clientReflector = clientReflector;
-        this.recordingInterceptor = new RequestRecordingInterceptor();
-    }
-
-    /**
-     * @return LoggedRequest that wire mock captured.
-     */
-    private static LoggedRequest getLoggedRequest() {
-        List<LoggedRequest> requests = WireMockUtils.findAllLoggedRequests();
-        assertEquals(1, requests.size());
-        return requests.get(0);
+        this.localhostOnlyForWiremockInterceptor = new LocalhostOnlyForWiremockInterceptor();
     }
 
     void runTest(TestCase testCase) throws Exception {
+        resetWireMock();
         ShapeModelReflector shapeModelReflector = createShapeModelReflector(testCase);
         AwsRequest request = createRequest(testCase, shapeModelReflector);
 
-        try {
-            if (!model.getShapes().get(testCase.getWhen().getOperationName() + "Request").isHasStreamingMember()) {
-                clientReflector.invokeMethod(testCase, request);
-            } else {
-                clientReflector.invokeMethod(testCase,
-                                             request,
-                                             RequestBody.fromString(shapeModelReflector.getStreamingMemberValue()));
-            }
-            Assert.fail("Expected SDK client to intercept and record request before transmission.");
-        } catch (InvocationTargetException e) {
-            if (e.getTargetException() instanceof StopExecutionException) {
-                SdkHttpRequest recordedRequest = recordingInterceptor.getRequest();
-                testCase.getThen().getMarshallingAssertion().assertMatches(getLoggedRequest());
-            } else {
-                throw e;
-            }
+        if (!model.getShapes().get(testCase.getWhen().getOperationName() + "Request").isHasStreamingMember()) {
+            clientReflector.invokeMethod(testCase, request);
+        } else {
+            clientReflector.invokeMethod(testCase,
+                                         request,
+                                         RequestBody.fromString(shapeModelReflector.getStreamingMemberValue()));
         }
+        testCase.getThen().getMarshallingAssertion()
+                .assertMatches(localhostOnlyForWiremockInterceptor.getLoggedRequestWithOriginalHost());
+    }
+
+    /**
+     * Reset wire mock and re-configure stubbing.
+     */
+    private void resetWireMock() {
+        WireMock.reset();
+        // Stub to return 200 for all requests
+        ResponseDefinitionBuilder responseDefBuilder = aResponse().withStatus(200);
+        // XML Unmarshallers expect at least one level in the XML document.
+        if (model.getMetadata().isXmlProtocol()) {
+            responseDefBuilder.withBody("<foo></foo>");
+        }
+        stubFor(any(urlMatching(".*")).willReturn(responseDefBuilder));
     }
 
     private AwsRequest createRequest(TestCase testCase, ShapeModelReflector shapeModelReflector) {
         return ((AwsRequest) shapeModelReflector.createShapeObject())
                 .toBuilder()
                 .overrideConfiguration(requestConfiguration -> requestConfiguration
                     .addPlugin(config -> {
-                        config.overrideConfiguration(c -> c.addExecutionInterceptor(recordingInterceptor));
+                        config.overrideConfiguration(c -> c.addExecutionInterceptor(localhostOnlyForWiremockInterceptor));
 
                         if (StringUtils.isNotBlank(testCase.getGiven().getHost())) {
                             config.endpointOverride(URI.create("https://" + testCase.getGiven().getHost()));
@@ -115,26 +116,49 @@ private String getOperationRequestClassName(String operationName) {
         return operationName + "Request";
     }
 
-    private static final class RequestRecordingInterceptor implements ExecutionInterceptor {
-        private SdkHttpRequest request;
+    /**
+     * Wiremock requires that requests use "localhost" as the host - any prefixes such as "foo.localhost" will
+     * result in a DNS lookup that will fail.  This interceptor modifies the request to ensure this and captures
+     * the original host.
+     */
+    private static final class LocalhostOnlyForWiremockInterceptor implements ExecutionInterceptor {
+        private String originalHost;
+        private String originalProtocol;
+        private int originalPort;
 
         @Override
-        public void beforeTransmission(Context.BeforeTransmission context, ExecutionAttributes executionAttributes) {
-            request = context.httpRequest();
-
-            // Log or record the request here
-            System.out.println("Recording Request:");
-            System.out.println("HTTP Method: " + request.method());
-            System.out.println("Endpoint: " + request.getUri());
-            System.out.println("Headers: " + request.headers());
-
-            throw new StopExecutionException();
+        public SdkHttpRequest modifyHttpRequest(Context.ModifyHttpRequest context, ExecutionAttributes executionAttributes) {
+            originalHost = context.httpRequest().host();
+            originalProtocol = context.httpRequest().protocol();
+            originalPort = context.httpRequest().port();
+
+            return context.httpRequest().toBuilder()
+                          .host("localhost")
+                          .port(WireMockUtils.port())
+                          .protocol("http")
+                          .build();
         }
 
-        public SdkHttpRequest getRequest() {
-            return request;
+        /**
+         * @return LoggedRequest that wire mock captured modified with the original host captured by this
+         * interceptor.
+         */
+        public LoggedRequest getLoggedRequestWithOriginalHost() {
+            List<LoggedRequest> requests = WireMockUtils.findAllLoggedRequests();
+            assertEquals(1, requests.size());
+            LoggedRequest loggedRequest = requests.get(0);
+            return new LoggedRequest(
+                loggedRequest.getUrl(),
+                originalProtocol + "://" + originalHost + ":" + originalPort,
+                loggedRequest.getMethod(),
+                loggedRequest.getClientIp(),
+                loggedRequest.getHeaders(),
+                loggedRequest.getCookies(),
+                loggedRequest.isBrowserProxyRequest(),
+                loggedRequest.getLoggedDate(),
+                loggedRequest.getBody(),
+                loggedRequest.getParts()
+            );
         }
     }
-
-    private static class StopExecutionException extends RuntimeException {}
 }
diff --git a/test/protocol-tests-core/src/main/resources/software/amazon/awssdk/protocol/suites/cases/query-input.json b/test/protocol-tests-core/src/main/resources/software/amazon/awssdk/protocol/suites/cases/query-input.json
@@ -28,7 +28,8 @@
     "id": "AwsQueryEndpointTrait",
     "description": "Operations can prepend to the given host if they define the\nendpoint trait.",
     "given": {
-      "input": {}
+      "input": {},
+      "host": "example.com"
     },
     "when": {
       "action": "marshall",
@@ -45,7 +46,8 @@
         "body": {
           "queryEquals": "Action=EndpointOperation&Version=2020-01-08"
         },
-        "uri": "/"
+        "uri": "/",
+        "resolvedHost": "foo.example.com"
       }
     }
   },
@@ -55,7 +57,8 @@
     "given": {
       "input": {
         "label": "bar"
-      }
+      },
+      "host": "example.com"
     },
     "when": {
       "action": "marshall",
@@ -72,27 +75,8 @@
         "body": {
           "queryEquals": "Action=EndpointWithHostLabelOperation&Version=2020-01-08&label=bar"
         },
-        "uri": "/"
-      }
-    }
-  },
-  {
-    "id": "QueryHostWithPath",
-    "description": "Custom endpoints supplied by users can have paths",
-    "given": {
-      "input": {}
-    },
-    "when": {
-      "action": "marshall",
-      "operation": "HostWithPathOperation"
-    },
-    "then": {
-      "serializedAs": {
-        "method": "POST",
-        "body": {
-          "queryEquals": "Action=HostWithPathOperation&Version=2020-01-08"
-        },
-        "uri": "/custom/"
+        "uri": "/",
+        "resolvedHost": "foo.bar.example.com"
       }
     }
   },
@@ -108,8 +92,7 @@
             "StringArg": "baz"
           }
         }
-      },
-      "host": "example.com"
+      }
     },
     "when": {
       "action": "marshall",
diff --git a/test/protocol-tests/src/main/resources/codegen-resources/smithy-query/customization.config b/test/protocol-tests/src/main/resources/codegen-resources/smithy-query/customization.config
@@ -1,3 +1,4 @@
 {
-    "skipEndpointTestGeneration": true
+    "skipEndpointTestGeneration": true,
+    "enableGenerateCompiledEndpointRules": true
 }

Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,4 @@`
`1`	`1`	`{`
`2`		`- "skipEndpointTestGeneration": true`
	`2`	`+ "skipEndpointTestGeneration": true,`
	`3`	`+ "enableGenerateCompiledEndpointRules": true`
`3`	`4`	`}`