add proxy wrapper to remove excess code

Author: jaydeluca

instrumentation/aws-sdk/aws-sdk-2.2/testing/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/AbstractAws2ClientCoreTest.java

@@ -31,7 +31,10 @@
 import io.opentelemetry.testing.internal.armeria.common.MediaType;
 import io.opentelemetry.testing.internal.armeria.testing.junit5.server.mock.MockWebServerExtension;
 import io.opentelemetry.testing.internal.armeria.testing.junit5.server.mock.RecordedRequest;
+import java.lang.reflect.Method;
+import java.lang.reflect.Proxy;
 import java.util.Arrays;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
 import java.util.function.Function;
@@ -243,23 +246,30 @@ private static void assertDynamoDbRequest(SpanDataAssert span, String operation)
             equalTo(maybeStable(DB_OPERATION), operation));
   }
 
+  @SuppressWarnings("unchecked")
+  protected static <T, U> T wrapClient(
+      Class<T> syncClientClass, Class<U> asyncClientClass, U asyncClient) {
+    return (T)
+        Proxy.newProxyInstance(
+            AbstractAws2ClientCoreTest.class.getClassLoader(),
+            new Class<?>[] {syncClientClass},
+            (proxy, method, args) -> {
+              Method asyncMethod =
+                  asyncClientClass.getMethod(method.getName(), method.getParameterTypes());
+              CompletableFuture<?> future =
+                  (CompletableFuture<?>) asyncMethod.invoke(asyncClient, args);
+              return future.get();
+            });
+  }
+
   private static Stream<Arguments> provideArguments() {
     return Stream.of(
         Arguments.of(
             "CreateTable",
-            (Function<DynamoDbClient, Object>) c -> c.createTable(createTableRequest()),
-            (Function<DynamoDbAsyncClient, Object>) c -> c.createTable(createTableRequest())),
+            (Function<DynamoDbClient, Object>) c -> c.createTable(createTableRequest())),
         Arguments.of(
             "DeleteItem",
             (Function<DynamoDbClient, Object>)
-                c ->
-                    c.deleteItem(
-                        DeleteItemRequest.builder()
-                            .tableName("sometable")
-                            .key(createTableRequestKey)
-                            .conditionExpression("property in (:one, :two)")
-                            .build()),
-            (Function<DynamoDbAsyncClient, Object>)
                 c ->
                     c.deleteItem(
                         DeleteItemRequest.builder()
@@ -270,20 +280,10 @@ private static Stream<Arguments> provideArguments() {
         Arguments.of(
             "DeleteTable",
             (Function<DynamoDbClient, Object>)
-                c -> c.deleteTable(DeleteTableRequest.builder().tableName("sometable").build()),
-            (Function<DynamoDbAsyncClient, Object>)
                 c -> c.deleteTable(DeleteTableRequest.builder().tableName("sometable").build())),
         Arguments.of(
             "GetItem",
             (Function<DynamoDbClient, Object>)
-                c ->
-                    c.getItem(
-                        GetItemRequest.builder()
-                            .tableName("sometable")
-                            .key(getItemRequestKey)
-                            .attributesToGet("propertyOne", "propertyTwo")
-                            .build()),
-            (Function<DynamoDbAsyncClient, Object>)
                 c ->
                     c.getItem(
                         GetItemRequest.builder()
@@ -294,14 +294,6 @@ private static Stream<Arguments> provideArguments() {
         Arguments.of(
             "PutItem",
             (Function<DynamoDbClient, Object>)
-                c ->
-                    c.putItem(
-                        PutItemRequest.builder()
-                            .tableName("sometable")
-                            .item(putItemRequestKey)
-                            .conditionExpression("attributeOne <> :someVal")
-                            .build()),
-            (Function<DynamoDbAsyncClient, Object>)
                 c ->
                     c.putItem(
                         PutItemRequest.builder()
@@ -312,16 +304,6 @@ private static Stream<Arguments> provideArguments() {
         Arguments.of(
             "Query",
             (Function<DynamoDbClient, Object>)
-                c ->
-                    c.query(
-                        QueryRequest.builder()
-                            .tableName("sometable")
-                            .select("ALL_ATTRIBUTES")
-                            .keyConditionExpression("attribute = :aValue")
-                            .filterExpression("anotherAttribute = :someVal")
-                            .limit(10)
-                            .build()),
-            (Function<DynamoDbAsyncClient, Object>)
                 c ->
                     c.query(
                         QueryRequest.builder()
@@ -334,20 +316,6 @@ private static Stream<Arguments> provideArguments() {
         Arguments.of(
             "UpdateItem",
             (Function<DynamoDbClient, Object>)
-                c ->
-                    c.updateItem(
-                        UpdateItemRequest.builder()
-                            .tableName("sometable")
-                            .key(
-                                ImmutableMap.of(
-                                    "keyOne",
-                                    AttributeValue.builder().s("value").build(),
-                                    "keyTwo",
-                                    AttributeValue.builder().s("differentValue").build()))
-                            .conditionExpression("attributeOne <> :someVal")
-                            .updateExpression("set attributeOne = :updateValue")
-                            .build()),
-            (Function<DynamoDbAsyncClient, Object>)
                 c ->
                     c.updateItem(
                         UpdateItemRequest.builder()
@@ -384,9 +352,7 @@ void testSendDynamoDbRequestWithBuilderAndMockedResponse(
   @ParameterizedTest
   @MethodSource("provideArguments")
   void testSendDynamoDbAsyncRequestWithBuilderAndMockedResponse(
-      String operation,
-      Function<DynamoDbClient, Object> call,
-      Function<DynamoDbAsyncClient, Object> asyncCall)
+      String operation, Function<DynamoDbClient, Object> call)
       throws ExecutionException, InterruptedException {
     DynamoDbAsyncClientBuilder builder = DynamoDbAsyncClient.builder();
     configureSdkClient(builder);
@@ -397,7 +363,8 @@ void testSendDynamoDbAsyncRequestWithBuilderAndMockedResponse(
             .credentialsProvider(CREDENTIALS_PROVIDER)
             .build();
     server.enqueue(HttpResponse.of(HttpStatus.OK, MediaType.PLAIN_TEXT_UTF_8, ""));
-    Object response = asyncCall.apply(client);
+    Object response =
+        call.apply(wrapClient(DynamoDbClient.class, DynamoDbAsyncClient.class, client));
     validateOperationResponse(operation, response);
   }
 }

instrumentation/aws-sdk/aws-sdk-2.2/testing/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/AbstractAws2ClientTest.java

@@ -92,12 +92,6 @@
 public abstract class AbstractAws2ClientTest extends AbstractAws2ClientCoreTest {
   private static final String QUEUE_URL = "http://xxx/somequeue";
 
-  private static void assumeSupportedConfig(String operation) {
-    Assumptions.assumeFalse(
-        operation.equals("SendMessage") && isSqsAttributeInjectionEnabled(),
-        "Cannot check Sqs.SendMessage here due to hard-coded MD5.");
-  }
-
   // Force localhost instead of relying on mock server because using ip is yet another corner case
   // of the virtual bucket changes introduced by aws sdk v2.18.0. When using IP, there is no way to
   // prefix the hostname with the bucket name as label.
@@ -115,6 +109,12 @@ private static void assumeSupportedConfig(String operation) {
           + " <ResponseMetadata><RequestId>0ac9cda2-bbf4-11d3-f92b-31fa5e8dbc99</RequestId></ResponseMetadata>"
           + "</DeleteOptionGroupResponse>";
 
+  private static void assumeSupportedConfig(String operation) {
+    Assumptions.assumeFalse(
+        operation.equals("SendMessage") && isSqsAttributeInjectionEnabled(),
+        "Cannot check Sqs.SendMessage here due to hard-coded MD5.");
+  }
+
   @SuppressWarnings("deprecation") // uses deprecated semconv
   private void clientAssertions(
       String service, String operation, String method, Object response, String requestId) {
@@ -322,7 +322,7 @@ void testS3AsyncSendOperationRequestWithBuilder(
   }
 
   @Test
-  void testKinesisSendOperationRequestWithBuilder() throws Exception {
+  void testKinesisSendOperationRequestWithBuilder() {
     KinesisClientBuilder builder = KinesisClient.builder();
     configureSdkClient(builder);
     KinesisClient client =
@@ -372,8 +372,6 @@ private static Stream<Arguments> provideSqsArguments() {
                   return HttpResponse.of(headers, HttpData.of(StandardCharsets.UTF_8, content));
                 },
             (Function<SqsClient, Object>)
-                c -> c.createQueue(CreateQueueRequest.builder().queueName("somequeue").build()),
-            (Function<SqsAsyncClient, Future<?>>)
                 c -> c.createQueue(CreateQueueRequest.builder().queueName("somequeue").build())),
         Arguments.of(
             "SendMessage",
@@ -407,10 +405,6 @@ private static Stream<Arguments> provideSqsArguments() {
                   return HttpResponse.of(headers, HttpData.of(StandardCharsets.UTF_8, content));
                 },
             (Function<SqsClient, Object>)
-                c ->
-                    c.sendMessage(
-                        SendMessageRequest.builder().queueUrl(QUEUE_URL).messageBody("").build()),
-            (Function<SqsAsyncClient, Future<?>>)
                 c ->
                     c.sendMessage(
                         SendMessageRequest.builder().queueUrl(QUEUE_URL).messageBody("").build())));
@@ -451,8 +445,7 @@ void testSqsAsyncSendOperationRequestWithBuilder(
       String operation,
       String requestId,
       Callable<HttpResponse> serverResponse,
-      Function<SqsClient, Object> call,
-      Function<SqsAsyncClient, Future<?>> asyncCall)
+      Function<SqsClient, Object> call)
       throws Exception {
     assumeSupportedConfig(operation);
 
@@ -466,8 +459,7 @@ void testSqsAsyncSendOperationRequestWithBuilder(
             .build();
 
     server.enqueue(serverResponse.call());
-    Future<?> response = asyncCall.apply(client);
-    response.get();
+    Object response = call.apply(wrapClient(SqsClient.class, SqsAsyncClient.class, client));
 
     clientAssertions("Sqs", operation, "POST", response, requestId);
   }
@@ -524,8 +516,7 @@ void testSnsSendOperationRequestWithBuilder(Function<SnsClient, Object> call) {
     clientAssertions("Sns", "Publish", "POST", response, "d74b8436-ae13-5ab4-a9ff-ce54dfea72a0");
   }
 
-  @ParameterizedTest
-  @MethodSource("provideSnsArguments")
+  @Test
   void testSnsAsyncSendOperationRequestWithBuilder() {
     SnsAsyncClientBuilder builder = SnsAsyncClient.builder();
     configureSdkClient(builder);
@@ -553,7 +544,7 @@ void testSnsAsyncSendOperationRequestWithBuilder() {
   }
 
   @Test
-  void testEc2SendOperationRequestWithBuilder() throws Exception {
+  void testEc2SendOperationRequestWithBuilder() {
     Ec2ClientBuilder builder = Ec2Client.builder();
     configureSdkClient(builder);
     Ec2Client client =
@@ -575,7 +566,7 @@ void testEc2SendOperationRequestWithBuilder() throws Exception {
   }
 
   @Test
-  void testEc2AsyncSendOperationRequestWithBuilder() throws Exception {
+  void testEc2AsyncSendOperationRequestWithBuilder() {
     Ec2AsyncClientBuilder builder = Ec2AsyncClient.builder();
     configureSdkClient(builder);
     Ec2AsyncClient client =
@@ -593,7 +584,7 @@ void testEc2AsyncSendOperationRequestWithBuilder() throws Exception {
   }
 
   @Test
-  void testRdsSendOperationRequestWithBuilder() throws Exception {
+  void testRdsSendOperationRequestWithBuilder() {
     RdsClientBuilder builder = RdsClient.builder();
     configureSdkClient(builder);
     RdsClient client =
@@ -636,7 +627,7 @@ void testRdsAsyncSendOperationRequestWithBuilder() {
   // spans because of https://github.com/aws/aws-sdk-java-v2/issues/1741. We should at least tweak
   // the instrumentation to add Events for retries instead.
   @Test
-  void testTimeoutAndRetryErrorsAreNotCaptured() throws Exception {
+  void testTimeoutAndRetryErrorsAreNotCaptured() {
     // One retry so two requests.
     server.enqueue(HttpResponse.delayed(HttpResponse.of(HttpStatus.OK), Duration.ofSeconds(5)));
     server.enqueue(HttpResponse.delayed(HttpResponse.of(HttpStatus.OK), Duration.ofSeconds(5)));