added unit tests

Author: liustve

awsagentprovider/src/test/java/software/amazon/opentelemetry/javaagent/providers/OtlpAwsSpanExporterTest.java

@@ -0,0 +1,174 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates.
+ *
+ * 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.opentelemetry.javaagent.providers;
+
+import static org.junit.jupiter.api.Assertions.*;
+import static org.mockito.Mockito.*;
+import static org.mockito.Mockito.when;
+
+import io.opentelemetry.exporter.otlp.http.trace.OtlpHttpSpanExporter;
+import io.opentelemetry.exporter.otlp.http.trace.OtlpHttpSpanExporterBuilder;
+import io.opentelemetry.sdk.common.CompletableResultCode;
+import io.opentelemetry.sdk.trace.export.SpanExporter;
+import java.net.URI;
+import java.util.List;
+import java.util.Map;
+import java.util.function.Consumer;
+import java.util.function.Supplier;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.MockedStatic;
+import org.mockito.junit.jupiter.MockitoExtension;
+import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
+import software.amazon.awssdk.auth.credentials.AwsCredentials;
+import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider;
+import software.amazon.awssdk.core.exception.SdkClientException;
+import software.amazon.awssdk.http.SdkHttpFullRequest;
+import software.amazon.awssdk.http.SdkHttpMethod;
+import software.amazon.awssdk.http.auth.aws.signer.AwsV4HttpSigner;
+import software.amazon.awssdk.http.auth.spi.signer.SignRequest.Builder;
+import software.amazon.awssdk.http.auth.spi.signer.SignedRequest;
+import software.amazon.awssdk.identity.spi.AwsCredentialsIdentity;
+
+@ExtendWith(MockitoExtension.class)
+public class OtlpAwsSpanExporterTest {
+  private static final String OTLP_CW_ENDPOINT = "https://xray.us-east-1.amazonaws.com/v1/traces";
+  private static final String AUTHORIZATION_HEADER = "Authorization";
+  private static final String X_AMZ_DATE_HEADER = "X-Amz-Date";
+  private static final String X_AMZ_SECURITY_TOKEN_HEADER = "X-Amz-Security-Token";
+
+  private static final String EXPECTED_AUTH_HEADER =
+      "AWS4-HMAC-SHA256 Credential=test_key/some_date/us-east-1/xray/aws4_request";
+  private static final String EXPECTED_AUTH_X_AMZ_DATE = "some_date";
+  private static final String EXPECTED_AUTH_SECURITY_TOKEN = "test_token";
+
+  AwsCredentials credentials = AwsBasicCredentials.create("test_access_key", "test_secret_key");
+  SignedRequest signedRequest =
+      SignedRequest.builder()
+          .request(
+              SdkHttpFullRequest.builder()
+                  .method(SdkHttpMethod.POST)
+                  .uri(URI.create(OTLP_CW_ENDPOINT))
+                  .putHeader(AUTHORIZATION_HEADER, EXPECTED_AUTH_HEADER)
+                  .putHeader("X-Amz-Date", EXPECTED_AUTH_X_AMZ_DATE)
+                  .putHeader("X-Amz-Security-Token", EXPECTED_AUTH_SECURITY_TOKEN)
+                  .build())
+          .build();
+
+  private MockedStatic<DefaultCredentialsProvider> mockDefaultCredentialsProvider;
+  private MockedStatic<AwsV4HttpSigner> mockAwsV4HttpSigner;
+  private MockedStatic<OtlpHttpSpanExporter> otlpSpanExporterMock;
+
+  @Mock private DefaultCredentialsProvider credentialsProvider;
+  @Mock private AwsV4HttpSigner signer;
+  @Mock private OtlpHttpSpanExporterBuilder mockBuilder;
+  @Mock private OtlpHttpSpanExporter mockExporter;
+
+  private ArgumentCaptor<Supplier<Map<String, String>>> headersCaptor;
+
+  @BeforeEach
+  void setup() {
+    this.mockDefaultCredentialsProvider = mockStatic(DefaultCredentialsProvider.class);
+    this.mockDefaultCredentialsProvider
+        .when(DefaultCredentialsProvider::create)
+        .thenReturn(credentialsProvider);
+
+    this.mockAwsV4HttpSigner = mockStatic(AwsV4HttpSigner.class);
+    this.mockAwsV4HttpSigner.when(AwsV4HttpSigner::create).thenReturn(this.signer);
+
+    this.otlpSpanExporterMock = mockStatic(OtlpHttpSpanExporter.class);
+
+    this.headersCaptor = ArgumentCaptor.forClass(Supplier.class);
+
+    when(OtlpHttpSpanExporter.builder()).thenReturn(mockBuilder);
+    when(this.mockBuilder.setEndpoint(any())).thenReturn(mockBuilder);
+    when(this.mockBuilder.setHeaders(headersCaptor.capture())).thenReturn(mockBuilder);
+    when(this.mockBuilder.build()).thenReturn(mockExporter);
+    when(this.mockExporter.export(any())).thenReturn(CompletableResultCode.ofSuccess());
+  }
+
+  @AfterEach
+  void afterEach() {
+    reset(this.signer, this.credentialsProvider);
+    this.mockDefaultCredentialsProvider.close();
+    this.mockAwsV4HttpSigner.close();
+    this.otlpSpanExporterMock.close();
+  }
+
+  @Test
+  void testAwsSpanExporterAddsSigV4Headers() {
+
+    when(this.credentialsProvider.resolveCredentials()).thenReturn(this.credentials);
+    when(this.signer.sign((Consumer<Builder<AwsCredentialsIdentity>>) any()))
+        .thenReturn(this.signedRequest);
+
+    SpanExporter exporter = new OtlpAwsSpanExporter(OTLP_CW_ENDPOINT);
+
+    exporter.export(List.of());
+
+    Map<String, String> headers = this.headersCaptor.getValue().get();
+
+    assertTrue(headers.containsKey(X_AMZ_DATE_HEADER));
+    assertTrue(headers.containsKey(AUTHORIZATION_HEADER));
+    assertTrue(headers.containsKey(X_AMZ_SECURITY_TOKEN_HEADER));
+
+    assertEquals(EXPECTED_AUTH_HEADER, headers.get(AUTHORIZATION_HEADER));
+    assertEquals(EXPECTED_AUTH_X_AMZ_DATE, headers.get(X_AMZ_DATE_HEADER));
+    assertEquals(EXPECTED_AUTH_SECURITY_TOKEN, headers.get(X_AMZ_SECURITY_TOKEN_HEADER));
+  }
+
+  @Test
+  void testAwsSpanExporterDoesNotAddSigV4HeadersIfFailureToRetrieveCredentials() {
+
+    when(this.credentialsProvider.resolveCredentials())
+        .thenThrow(SdkClientException.builder().message("bad credentials").build());
+
+    SpanExporter exporter = new OtlpAwsSpanExporter(OTLP_CW_ENDPOINT);
+
+    exporter.export(List.of());
+
+    Supplier<Map<String, String>> headersSupplier = headersCaptor.getValue();
+    Map<String, String> headers = headersSupplier.get();
+
+    assertFalse(headers.containsKey(X_AMZ_DATE_HEADER));
+    assertFalse(headers.containsKey(AUTHORIZATION_HEADER));
+    assertFalse(headers.containsKey(X_AMZ_SECURITY_TOKEN_HEADER));
+
+    verifyNoInteractions(this.signer);
+  }
+
+  @Test
+  void testAwsSpanExporterDoesNotAddSigV4HeadersIfFailureToSignHeaders() {
+
+    when(this.credentialsProvider.resolveCredentials()).thenReturn(this.credentials);
+    when(this.signer.sign((Consumer<Builder<AwsCredentialsIdentity>>) any()))
+        .thenThrow(SdkClientException.builder().message("bad signature").build());
+
+    SpanExporter exporter = new OtlpAwsSpanExporter(OTLP_CW_ENDPOINT);
+
+    exporter.export(List.of());
+
+    Map<String, String> headers = this.headersCaptor.getValue().get();
+
+    assertFalse(headers.containsKey(X_AMZ_DATE_HEADER));
+    assertFalse(headers.containsKey(AUTHORIZATION_HEADER));
+    assertFalse(headers.containsKey(X_AMZ_SECURITY_TOKEN_HEADER));
+  }
+}