Add unit tests to verify quota project behavior

Author: psx95

gcp-auth-extension/build.gradle.kts

@@ -29,6 +29,7 @@ dependencies {
   testCompileOnly("com.google.auto.service:auto-service-annotations")
   testRuntimeOnly("org.junit.jupiter:junit-jupiter-engine")
   testImplementation("org.junit.jupiter:junit-jupiter-api")
+  testCompileOnly("org.junit.jupiter:junit-jupiter-params")
 
   testImplementation("io.opentelemetry:opentelemetry-api")
   testImplementation("io.opentelemetry:opentelemetry-exporter-otlp")
@@ -45,6 +46,9 @@ dependencies {
   testImplementation("org.springframework.boot:spring-boot-starter:2.7.18")
   testImplementation("org.springframework.boot:spring-boot-starter-test:2.7.18")
 
+  testAnnotationProcessor("com.google.auto.value:auto-value")
+  testCompileOnly("com.google.auto.value:auto-value-annotations")
+
   agent("io.opentelemetry.javaagent:opentelemetry-javaagent")
 }
 

gcp-auth-extension/src/test/java/io/opentelemetry/contrib/gcp/auth/GcpAuthAutoConfigurationCustomizerProviderTest.java

@@ -14,6 +14,7 @@
 
 import com.google.auth.oauth2.AccessToken;
 import com.google.auth.oauth2.GoogleCredentials;
+import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableMap;
 import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.api.trace.Span;
@@ -46,9 +47,14 @@
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Supplier;
+import java.util.stream.Stream;
+import javax.annotation.Nullable;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Captor;
 import org.mockito.Mock;
@@ -149,19 +155,11 @@ public void testCustomizerOtlpGrpc() {
     // Prepare mocks
     prepareMockBehaviorForGoogleCredentials();
     OtlpGrpcSpanExporter mockOtlpGrpcSpanExporter = Mockito.mock(OtlpGrpcSpanExporter.class);
-    OtlpGrpcSpanExporterBuilder otlpSpanExporterBuilder = OtlpGrpcSpanExporter.builder();
     OtlpGrpcSpanExporterBuilder spyOtlpGrpcSpanExporterBuilder =
-        Mockito.spy(otlpSpanExporterBuilder);
-    Mockito.when(spyOtlpGrpcSpanExporterBuilder.build()).thenReturn(mockOtlpGrpcSpanExporter);
-    Mockito.when(mockOtlpGrpcSpanExporter.shutdown()).thenReturn(CompletableResultCode.ofSuccess());
+        Mockito.spy(OtlpGrpcSpanExporter.builder());
     List<SpanData> exportedSpans = new ArrayList<>();
-    Mockito.when(mockOtlpGrpcSpanExporter.export(Mockito.anyCollection()))
-        .thenAnswer(
-            invocationOnMock -> {
-              exportedSpans.addAll(invocationOnMock.getArgument(0));
-              return CompletableResultCode.ofSuccess();
-            });
-    Mockito.when(mockOtlpGrpcSpanExporter.toBuilder()).thenReturn(spyOtlpGrpcSpanExporterBuilder);
+    configureGrpcMockExporters(
+        mockOtlpGrpcSpanExporter, spyOtlpGrpcSpanExporterBuilder, exportedSpans);
 
     try (MockedStatic<GoogleCredentials> googleCredentialsMockedStatic =
         Mockito.mockStatic(GoogleCredentials.class)) {
@@ -213,6 +211,186 @@ public void testCustomizerFailWithMissingResourceProject() {
     }
   }
 
+  @ParameterizedTest
+  @MethodSource("provideQuotaBehaviorTestCases")
+  @SuppressWarnings("CannotMockMethod")
+  public void testQuotaProjectBehavior(QuotaProjectIdTestBehavior testCase) {
+    // Set resource project system property
+    System.setProperty(
+        ConfigurableOption.GOOGLE_CLOUD_PROJECT.getSystemProperty(), DUMMY_GCP_RESOURCE_PROJECT_ID);
+    // Configure mock credentials to return fake access token
+    Mockito.when(mockedGoogleCredentials.getAccessToken())
+        .thenReturn(new AccessToken("fake", Date.from(Instant.now())));
+
+    // To prevent unncecessary stubbings, mock getQuotaProjectId only when necessary
+    if (testCase.getUserSpecifiedQuotaProjectId() == null
+        || testCase.getUserSpecifiedQuotaProjectId().isEmpty()) {
+      String quotaProjectFromCredential =
+          testCase.getIsQuotaProjectPresentInCredentials() ? DUMMY_GCP_QUOTA_PROJECT_ID : null;
+      Mockito.when(mockedGoogleCredentials.getQuotaProjectId())
+          .thenReturn(quotaProjectFromCredential);
+    }
+
+    // configure environment according to test case
+    String quotaProjectId = testCase.getUserSpecifiedQuotaProjectId(); // maybe empty string
+    if (testCase.getUserSpecifiedQuotaProjectId() != null) {
+      // user specified a quota project id
+      System.setProperty(
+          ConfigurableOption.GOOGLE_CLOUD_QUOTA_PROJECT.getSystemProperty(), quotaProjectId);
+    }
+
+    // prepare mock exporter
+    OtlpGrpcSpanExporter mockOtlpGrpcSpanExporter = Mockito.mock(OtlpGrpcSpanExporter.class);
+    OtlpGrpcSpanExporterBuilder spyOtlpGrpcSpanExporterBuilder =
+        Mockito.spy(OtlpGrpcSpanExporter.builder());
+    List<SpanData> exportedSpans = new ArrayList<>();
+    configureGrpcMockExporters(
+        mockOtlpGrpcSpanExporter, spyOtlpGrpcSpanExporterBuilder, exportedSpans);
+
+    try (MockedStatic<GoogleCredentials> googleCredentialsMockedStatic =
+        Mockito.mockStatic(GoogleCredentials.class)) {
+      googleCredentialsMockedStatic
+          .when(GoogleCredentials::getApplicationDefault)
+          .thenReturn(mockedGoogleCredentials);
+
+      // Export telemetry to capture headers in the export calls
+      OpenTelemetrySdk sdk = buildOpenTelemetrySdkWithExporter(mockOtlpGrpcSpanExporter);
+      generateTestSpan(sdk);
+      CompletableResultCode code = sdk.shutdown();
+      CompletableResultCode joinResult = code.join(10, TimeUnit.SECONDS);
+      assertTrue(joinResult.isSuccess());
+      Mockito.verify(spyOtlpGrpcSpanExporterBuilder, Mockito.times(1))
+          .setHeaders(headerSupplierCaptor.capture());
+
+      // assert that the Authorization bearer token header is present
+      Map<String, String> exportHeaders = headerSupplierCaptor.getValue().get();
+      assertThat(exportHeaders).containsEntry("Authorization", "Bearer fake");
+
+      if (testCase.getExpectedQuotaProjectInHeader() == null) {
+        // there should be no user quota project header
+        assertThat(exportHeaders).doesNotContainKey(QUOTA_USER_PROJECT_HEADER);
+      } else {
+        // there should be user quota project header with expected value
+        assertThat(exportHeaders)
+            .containsEntry(QUOTA_USER_PROJECT_HEADER, testCase.getExpectedQuotaProjectInHeader());
+      }
+    }
+  }
+
+  /**
+   * Test cases specifying expected value for the user quota project header given the user input and
+   * the current credentials state.
+   *
+   * <p>{@code null} for {@link QuotaProjectIdTestBehavior#getUserSpecifiedQuotaProjectId()}
+   * indicates the case of user not specifying the quota project ID.
+   *
+   * <p>{@code null} value for {@link QuotaProjectIdTestBehavior#getExpectedQuotaProjectInHeader()}
+   * indicates the expectation that the QUOTA_USER_PROJECT_HEADER should not be present in the
+   * export headers.
+   *
+   * <p>{@code true} for {@link QuotaProjectIdTestBehavior#getIsQuotaProjectPresentInCredentials()}
+   * indicates that the mocked credentials are configured to provide DUMMY_GCP_QUOTA_PROJECT_ID as
+   * the quota project ID.
+   */
+  private static Stream<Arguments> provideQuotaBehaviorTestCases() {
+    return Stream.of(
+        Arguments.of(
+            QuotaProjectIdTestBehavior.builder()
+                .setUserSpecifiedQuotaProjectId(DUMMY_GCP_QUOTA_PROJECT_ID)
+                .setIsQuotaProjectPresentInCredentials(true)
+                .setExpectedQuotaProjectInHeader(DUMMY_GCP_QUOTA_PROJECT_ID)
+                .build()),
+        Arguments.of(
+            QuotaProjectIdTestBehavior.builder()
+                .setUserSpecifiedQuotaProjectId(DUMMY_GCP_QUOTA_PROJECT_ID)
+                .setIsQuotaProjectPresentInCredentials(false)
+                .setExpectedQuotaProjectInHeader(DUMMY_GCP_QUOTA_PROJECT_ID)
+                .build()),
+        Arguments.of(
+            QuotaProjectIdTestBehavior.builder()
+                .setUserSpecifiedQuotaProjectId("my-custom-quota-project-id")
+                .setIsQuotaProjectPresentInCredentials(true)
+                .setExpectedQuotaProjectInHeader("my-custom-quota-project-id")
+                .build()),
+        Arguments.of(
+            QuotaProjectIdTestBehavior.builder()
+                .setUserSpecifiedQuotaProjectId("my-custom-quota-project-id")
+                .setIsQuotaProjectPresentInCredentials(false)
+                .setExpectedQuotaProjectInHeader("my-custom-quota-project-id")
+                .build()),
+        Arguments.of(
+            QuotaProjectIdTestBehavior.builder()
+                .setUserSpecifiedQuotaProjectId("") // user explicitly specifies empty
+                .setIsQuotaProjectPresentInCredentials(true)
+                .setExpectedQuotaProjectInHeader(DUMMY_GCP_QUOTA_PROJECT_ID)
+                .build()),
+        Arguments.of(
+            QuotaProjectIdTestBehavior.builder()
+                .setUserSpecifiedQuotaProjectId(null) // user omits specifying quota project
+                .setIsQuotaProjectPresentInCredentials(true)
+                .setExpectedQuotaProjectInHeader(DUMMY_GCP_QUOTA_PROJECT_ID)
+                .build()),
+        Arguments.of(
+            QuotaProjectIdTestBehavior.builder()
+                .setUserSpecifiedQuotaProjectId("")
+                .setIsQuotaProjectPresentInCredentials(false)
+                .setExpectedQuotaProjectInHeader(null)
+                .build()),
+        Arguments.of(
+            QuotaProjectIdTestBehavior.builder()
+                .setUserSpecifiedQuotaProjectId(null)
+                .setIsQuotaProjectPresentInCredentials(false)
+                .setExpectedQuotaProjectInHeader(null)
+                .build()));
+  }
+
+  // Configure necessary behavior on the Grpc mock exporters to work
+  // TODO: Potential improvement - make this work for Http exporter as well.
+  private static void configureGrpcMockExporters(
+      OtlpGrpcSpanExporter mockGrpcExporter,
+      OtlpGrpcSpanExporterBuilder spyGrpcExporterBuilder,
+      List<SpanData> exportedSpanContainer) {
+    Mockito.when(spyGrpcExporterBuilder.build()).thenReturn(mockGrpcExporter);
+    Mockito.when(mockGrpcExporter.shutdown()).thenReturn(CompletableResultCode.ofSuccess());
+    Mockito.when(mockGrpcExporter.toBuilder()).thenReturn(spyGrpcExporterBuilder);
+    Mockito.when(mockGrpcExporter.export(Mockito.anyCollection()))
+        .thenAnswer(
+            invocationOnMock -> {
+              exportedSpanContainer.addAll(invocationOnMock.getArgument(0));
+              return CompletableResultCode.ofSuccess();
+            });
+  }
+
+  @AutoValue
+  abstract static class QuotaProjectIdTestBehavior {
+    // A null user specified quota represents the use case where user omits specifying quota
+    @Nullable
+    abstract String getUserSpecifiedQuotaProjectId();
+
+    abstract boolean getIsQuotaProjectPresentInCredentials();
+
+    // If expected quota project in header is null, the header entry should not be present in export
+    @Nullable
+    abstract String getExpectedQuotaProjectInHeader();
+
+    static Builder builder() {
+      return new AutoValue_GcpAuthAutoConfigurationCustomizerProviderTest_QuotaProjectIdTestBehavior
+          .Builder();
+    }
+
+    @AutoValue.Builder
+    abstract static class Builder {
+      abstract Builder setUserSpecifiedQuotaProjectId(String quotaProjectId);
+
+      abstract Builder setIsQuotaProjectPresentInCredentials(
+          boolean quotaProjectPresentInCredentials);
+
+      abstract Builder setExpectedQuotaProjectInHeader(String expectedQuotaProjectInHeader);
+
+      abstract QuotaProjectIdTestBehavior build();
+    }
+  }
+
   @SuppressWarnings("CannotMockMethod")
   private void prepareMockBehaviorForGoogleCredentials() {
     Mockito.when(mockedGoogleCredentials.getQuotaProjectId())
@@ -256,7 +434,9 @@ public String getName() {
   private static boolean authHeadersQuotaProjectIsPresent(Map<String, String> headers) {
     Set<Entry<String, String>> headerEntrySet = headers.entrySet();
     return headerEntrySet.contains(
-            new SimpleEntry<>(QUOTA_USER_PROJECT_HEADER, DUMMY_GCP_QUOTA_PROJECT_ID))
+            new SimpleEntry<>(
+                QUOTA_USER_PROJECT_HEADER,
+                GcpAuthAutoConfigurationCustomizerProviderTest.DUMMY_GCP_QUOTA_PROJECT_ID))
         && headerEntrySet.contains(new SimpleEntry<>("Authorization", "Bearer fake"));
   }