Async core metrics part3: add tests to verify async core metrics for event streaming apis

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/util/MetricUtils.java

@@ -16,6 +16,7 @@
 package software.amazon.awssdk.core.internal.util;
 
 import static software.amazon.awssdk.core.client.config.SdkClientOption.METRIC_PUBLISHER;
+import static software.amazon.awssdk.core.http.HttpResponseHandler.X_AMZN_REQUEST_ID_HEADER;
 
 import java.time.Duration;
 import java.util.Optional;
@@ -106,6 +107,8 @@ public static void collectHttpMetrics(MetricCollector metricCollector, SdkHttpFu
         metricCollector.reportMetric(CoreMetric.HTTP_STATUS_CODE, httpResponse.statusCode());
         httpResponse.firstMatchingHeader("x-amz-request-id")
                     .ifPresent(v -> metricCollector.reportMetric(CoreMetric.AWS_REQUEST_ID, v));
+        httpResponse.firstMatchingHeader(X_AMZN_REQUEST_ID_HEADER)
+                    .ifPresent(v -> metricCollector.reportMetric(CoreMetric.AWS_REQUEST_ID, v));
         httpResponse.firstMatchingHeader("x-amz-id-2")
                     .ifPresent(v -> metricCollector.reportMetric(CoreMetric.AWS_EXTENDED_REQUEST_ID, v));
     }

services/transcribestreaming/src/it/java/software/amazon/awssdk/services/transcribestreaming/TranscribeStreamingIntegrationTest.java

@@ -16,18 +16,21 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
 import static software.amazon.awssdk.http.Header.CONTENT_TYPE;
 
 import java.io.File;
 import java.io.FileInputStream;
 import java.io.FileNotFoundException;
 import java.io.InputStream;
-import java.net.URISyntaxException;
+import java.time.Duration;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import org.junit.BeforeClass;
 import org.junit.Test;
+import org.mockito.ArgumentCaptor;
 import org.reactivestreams.Publisher;
 import org.reactivestreams.Subscriber;
 import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
@@ -36,6 +39,7 @@
 import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
 import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
 import software.amazon.awssdk.core.internal.util.Mimetype;
+import software.amazon.awssdk.core.metrics.CoreMetric;
 import software.amazon.awssdk.metrics.MetricCollection;
 import software.amazon.awssdk.metrics.MetricPublisher;
 import software.amazon.awssdk.regions.Region;
@@ -44,6 +48,7 @@
 import software.amazon.awssdk.services.transcribestreaming.model.MediaEncoding;
 import software.amazon.awssdk.services.transcribestreaming.model.StartStreamTranscriptionRequest;
 import software.amazon.awssdk.services.transcribestreaming.model.StartStreamTranscriptionResponseHandler;
+import software.amazon.awssdk.utils.Logger;
 
 /**
  * An example test class to show the usage of
@@ -53,43 +58,36 @@
  * The audio files used in this class don't have voice, so there won't be any transcripted text would be empty
  */
 public class TranscribeStreamingIntegrationTest {
+    private static final Logger log = Logger.loggerFor(TranscribeStreamingIntegrationTest.class);
 
     private static TranscribeStreamingAsyncClient client;
 
+    private static MetricPublisher mockPublisher;
+
     @BeforeClass
-    public static void setup() throws URISyntaxException {
+    public static void setup() {
+        mockPublisher = mock(MetricPublisher.class);
         client = TranscribeStreamingAsyncClient.builder()
                                                .region(Region.US_EAST_1)
                                                .overrideConfiguration(b -> b.addExecutionInterceptor(new VerifyHeaderInterceptor())
-                                               .metricPublisher(new MetricPublisher() {
-                                                   @Override
-                                                   public void publish(MetricCollection metricCollection) {
-                                                       System.out.println(metricCollection);
-                                                   }
-
-                                                   @Override
-                                                   public void close() {
-
-                                                   }
-                                               })
-                                               )
+                                               .metricPublisher(mockPublisher))
                                                .credentialsProvider(getCredentials())
                                                .build();
     }
 
     @Test
-    public void testFileWith16kRate() throws ExecutionException, InterruptedException, URISyntaxException {
+    public void testFileWith16kRate() throws InterruptedException {
         CompletableFuture<Void> result = client.startStreamTranscription(getRequest(16_000),
                                                                          new AudioStreamPublisher(
                                                                              getInputStream("silence_16kHz_s16le.wav")),
                                                                          TestResponseHandlers.responseHandlerBuilder_Classic());
 
-        // Blocking call to keep the main thread for shutting down
-        result.get();
+        result.join();
+        verifyMetrics();
     }
 
     @Test
-    public void testFileWith8kRate() throws ExecutionException, InterruptedException, URISyntaxException {
+    public void testFileWith8kRate() throws ExecutionException, InterruptedException {
         CompletableFuture<Void> result = client.startStreamTranscription(getRequest(8_000),
                                                                          new AudioStreamPublisher(
                                                                              getInputStream("silence_8kHz_s16le.wav")),
@@ -143,4 +141,34 @@ public void beforeTransmission(Context.BeforeTransmission context, ExecutionAttr
             assertThat(contentTypeHeader.get(0)).isEqualTo(Mimetype.MIMETYPE_EVENT_STREAM);
         }
     }
+
+    private void verifyMetrics() throws InterruptedException {
+        // wait for 100ms for metrics to be delivered to mockPublisher
+        Thread.sleep(100);
+        ArgumentCaptor<MetricCollection> collectionCaptor = ArgumentCaptor.forClass(MetricCollection.class);
+        verify(mockPublisher).publish(collectionCaptor.capture());
+        MetricCollection capturedCollection = collectionCaptor.getValue();
+        assertThat(capturedCollection.name()).isEqualTo("ApiCall");
+        log.info(() -> "captured collection: " + capturedCollection);
+
+        assertThat(capturedCollection.metricValues(CoreMetric.CREDENTIALS_FETCH_DURATION).get(0))
+            .isGreaterThanOrEqualTo(Duration.ZERO);
+        assertThat(capturedCollection.metricValues(CoreMetric.MARSHALLING_DURATION).get(0))
+            .isGreaterThanOrEqualTo(Duration.ZERO);
+        assertThat(capturedCollection.metricValues(CoreMetric.API_CALL_DURATION).get(0))
+            .isGreaterThan(Duration.ZERO);
+
+        MetricCollection attemptCollection = capturedCollection.children().get(0);
+        assertThat(attemptCollection.name()).isEqualTo("ApiCallAttempt");
+        assertThat(attemptCollection.children()).isEmpty();
+        assertThat(attemptCollection.metricValues(CoreMetric.HTTP_STATUS_CODE))
+            .containsExactly(200);
+        assertThat(attemptCollection.metricValues(CoreMetric.SIGNING_DURATION).get(0))
+            .isGreaterThanOrEqualTo(Duration.ZERO);
+        assertThat(attemptCollection.metricValues(CoreMetric.AWS_REQUEST_ID).get(0)).isNotEmpty();
+
+        assertThat(attemptCollection.metricValues(CoreMetric.HTTP_REQUEST_ROUND_TRIP_TIME).get(0))
+            .isGreaterThanOrEqualTo(Duration.ofMillis(100));
+    }
+
 }

test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/metrics/async/AsyncEventStreamingCoreMetricsTest.java

@@ -0,0 +1,110 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.services.metrics.async;
+
+import static org.mockito.Mockito.when;
+
+import com.github.tomakehurst.wiremock.junit.WireMockRule;
+import java.net.URI;
+import java.util.concurrent.CompletableFuture;
+import java.util.function.Supplier;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.runners.MockitoJUnitRunner;
+import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
+import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
+import software.amazon.awssdk.core.async.EmptyPublisher;
+import software.amazon.awssdk.core.signer.NoOpSigner;
+import software.amazon.awssdk.metrics.MetricPublisher;
+import software.amazon.awssdk.services.protocolrestjson.ProtocolRestJsonAsyncClient;
+import software.amazon.awssdk.services.protocolrestjson.model.EventStreamOperationRequest;
+import software.amazon.awssdk.services.protocolrestjson.model.EventStreamOperationResponseHandler;
+
+/**
+ * Core metrics test for async streaming API
+ */
+@RunWith(MockitoJUnitRunner.class)
+public class AsyncEventStreamingCoreMetricsTest extends BaseAsyncCoreMetricsTest {
+    @Rule
+    public WireMockRule wireMock = new WireMockRule(0);
+
+    @Mock
+    private AwsCredentialsProvider mockCredentialsProvider;
+
+    @Mock
+    private MetricPublisher mockPublisher;
+
+
+    private ProtocolRestJsonAsyncClient client;
+
+    @Before
+    public void setup() {
+        client = ProtocolRestJsonAsyncClient.builder()
+                                            .credentialsProvider(mockCredentialsProvider)
+                                            .endpointOverride(URI.create("http://localhost:" + wireMock.port()))
+                                            .overrideConfiguration(c -> c.metricPublisher(mockPublisher)
+                                                                         .retryPolicy(b -> b.numRetries(MAX_RETRIES)))
+                                            .build();
+
+        when(mockCredentialsProvider.resolveCredentials()).thenAnswer(invocation -> {
+            try {
+                Thread.sleep(100);
+            } catch (InterruptedException ie) {
+                ie.printStackTrace();
+            }
+            return AwsBasicCredentials.create("foo", "bar");
+        });
+    }
+
+    @After
+    public void teardown() {
+        wireMock.resetAll();
+        if (client != null) {
+            client.close();
+        }
+        client = null;
+    }
+
+    @Override
+    String operationName() {
+        return "EventStreamOperation";
+    }
+
+    @Override
+    Supplier<CompletableFuture<?>> callable() {
+        return () -> client.eventStreamOperation(EventStreamOperationRequest.builder().overrideConfiguration(b -> b.signer(new NoOpSigner())).build(),
+                                                 new EmptyPublisher<>(),
+                                                 EventStreamOperationResponseHandler.builder()
+                                                                                    .subscriber(b -> {})
+                                                                                    .build());
+    }
+
+    @Override
+    MetricPublisher publisher() {
+        return mockPublisher;
+    }
+
+    void addDelayIfNeeded() {
+        try {
+            Thread.sleep(100);
+        } catch (InterruptedException e) {
+            e.printStackTrace();
+        }
+    }
+}

test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/metrics/async/BaseAsyncCoreMetricsTest.java

@@ -50,8 +50,8 @@ public abstract class BaseAsyncCoreMetricsTest {
     @Test
     public void apiCall_operationSuccessful_addsMetrics() {
         stubSuccessfulResponse();
-
         callable().get().join();
+        addDelayIfNeeded();
 
         ArgumentCaptor<MetricCollection> collectionCaptor = ArgumentCaptor.forClass(MetricCollection.class);
         verify(publisher()).publish(collectionCaptor.capture());
@@ -74,6 +74,7 @@ public void apiCall_operationSuccessful_addsMetrics() {
     public void apiCall_allRetryAttemptsFailedOf500() {
         stubErrorResponse();
         assertThatThrownBy(() -> callable().get().join()).hasCauseInstanceOf(EmptyModeledException.class);
+        addDelayIfNeeded();
 
         ArgumentCaptor<MetricCollection> collectionCaptor = ArgumentCaptor.forClass(MetricCollection.class);
         verify(publisher()).publish(collectionCaptor.capture());
@@ -89,6 +90,7 @@ public void apiCall_allRetryAttemptsFailedOf500() {
     public void apiCall_allRetryAttemptsFailedOfNetworkError() {
         stubNetworkError();
         assertThatThrownBy(() -> callable().get().join()).hasCauseInstanceOf(SdkClientException.class);
+        addDelayIfNeeded();
 
         ArgumentCaptor<MetricCollection> collectionCaptor = ArgumentCaptor.forClass(MetricCollection.class);
         verify(publisher()).publish(collectionCaptor.capture());
@@ -115,6 +117,7 @@ public void apiCall_allRetryAttemptsFailedOfNetworkError() {
     public void apiCall_firstAttemptFailedRetrySucceeded() {
         stubSuccessfulRetry();
         callable().get().join();
+        addDelayIfNeeded();
 
         ArgumentCaptor<MetricCollection> collectionCaptor = ArgumentCaptor.forClass(MetricCollection.class);
         verify(publisher()).publish(collectionCaptor.capture());
@@ -130,6 +133,14 @@ public void apiCall_firstAttemptFailedRetrySucceeded() {
         verifySuccessfulApiCallAttemptCollection(successfulAttempt);
     }
 
+    /**
+     * Adds delay after calling CompletableFuture.join to wait for publisher to get metrics.
+     * no op by default, can be overridden by subclasses
+     */
+    void addDelayIfNeeded() {
+        // no op by default
+    }
+
     abstract String operationName();
 
     abstract Supplier<CompletableFuture<?>> callable();