Skip to content

Commit 2dd2e13

Browse files
committed
Add unit test for extension
1 parent 250b108 commit 2dd2e13

File tree

4 files changed

+268
-2
lines changed

4 files changed

+268
-2
lines changed

build.gradle

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -225,6 +225,7 @@ subprojects {
225225
junit5_runtime : "org.junit.jupiter:junit-jupiter-engine:${junit5Version}",
226226
junit5_params : "org.junit.jupiter:junit-jupiter-params:${junit5Version}",
227227
mockito : "org.mockito:mockito-inline:${mockitoVersion}",
228+
mockito_jupiter: "org.mockito:mockito-junit-jupiter:${mockitoVersion}",
228229
slf4j_simple: "org.slf4j:slf4j-simple:${slf4jVersion}",
229230
opentelemetry_sdk_testing: "io.opentelemetry:opentelemetry-sdk-testing:${openTelemetryVersion}",
230231
test_containers: "org.testcontainers:testcontainers:${testContainersVersion}",

javaagent-extensions/gcp-auth/build.gradle

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,9 +55,16 @@ dependencies {
5555

5656
// test dependencies
5757
testImplementation(testLibraries.junit5)
58+
testImplementation(libraries.opentelemetry_api)
59+
testImplementation(libraries.opentelemetry_otlp_exporter)
5860
testRuntimeOnly(testLibraries.junit5_runtime)
61+
testImplementation(testLibraries.opentelemetry_sdk_testing)
62+
testImplementation(testLibraries.mockito)
63+
testImplementation(testLibraries.mockito_jupiter)
64+
testImplementation(libraries.opentelemetry_sdk_autoconf)
5965
// OTel instrumentation used in the sample app to facilitate integration testing
6066
agent agentLibraries.agent
67+
testImplementation 'org.awaitility:awaitility:4.2.2'
6168
testImplementation 'org.apache.httpcomponents:httpclient:4.5.14'
6269
includeDeps 'org.apache.httpcomponents:httpclient:4.5.14'
6370
}

javaagent-extensions/gcp-auth/src/main/java/com/google/cloud/opentelemetry/extension/auth/GcpAuthAutoConfigurationCustomizerProvider.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -50,8 +50,8 @@
5050
public class GcpAuthAutoConfigurationCustomizerProvider
5151
implements AutoConfigurationCustomizerProvider {
5252

53-
private static final String QUOTA_USER_PROJECT_HEADER = "X-Goog-User-Project";
54-
private static final String GCP_USER_PROJECT_ID_KEY = "gcp.project_id";
53+
static final String QUOTA_USER_PROJECT_HEADER = "X-Goog-User-Project";
54+
static final String GCP_USER_PROJECT_ID_KEY = "gcp.project_id";
5555

5656
/**
5757
* Customizes the provided {@link AutoConfigurationCustomizer}.
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,258 @@
1+
/*
2+
* Copyright 2024 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.google.cloud.opentelemetry.extension.auth;
17+
18+
import static com.google.cloud.opentelemetry.extension.auth.GcpAuthAutoConfigurationCustomizerProvider.GCP_USER_PROJECT_ID_KEY;
19+
import static com.google.cloud.opentelemetry.extension.auth.GcpAuthAutoConfigurationCustomizerProvider.QUOTA_USER_PROJECT_HEADER;
20+
import static org.awaitility.Awaitility.await;
21+
import static org.junit.jupiter.api.Assertions.*;
22+
23+
import com.google.auth.oauth2.AccessToken;
24+
import com.google.auth.oauth2.GoogleCredentials;
25+
import com.google.common.collect.ImmutableMap;
26+
import io.opentelemetry.api.common.AttributeKey;
27+
import io.opentelemetry.api.trace.Span;
28+
import io.opentelemetry.context.Scope;
29+
import io.opentelemetry.exporter.otlp.http.trace.OtlpHttpSpanExporter;
30+
import io.opentelemetry.exporter.otlp.http.trace.OtlpHttpSpanExporterBuilder;
31+
import io.opentelemetry.exporter.otlp.trace.OtlpGrpcSpanExporter;
32+
import io.opentelemetry.exporter.otlp.trace.OtlpGrpcSpanExporterBuilder;
33+
import io.opentelemetry.sdk.OpenTelemetrySdk;
34+
import io.opentelemetry.sdk.autoconfigure.AutoConfiguredOpenTelemetrySdk;
35+
import io.opentelemetry.sdk.autoconfigure.AutoConfiguredOpenTelemetrySdkBuilder;
36+
import io.opentelemetry.sdk.autoconfigure.internal.AutoConfigureUtil;
37+
import io.opentelemetry.sdk.autoconfigure.internal.ComponentLoader;
38+
import io.opentelemetry.sdk.autoconfigure.internal.SpiHelper;
39+
import io.opentelemetry.sdk.autoconfigure.spi.ConfigProperties;
40+
import io.opentelemetry.sdk.autoconfigure.spi.traces.ConfigurableSpanExporterProvider;
41+
import io.opentelemetry.sdk.common.CompletableResultCode;
42+
import io.opentelemetry.sdk.trace.data.SpanData;
43+
import io.opentelemetry.sdk.trace.export.SpanExporter;
44+
import java.time.Duration;
45+
import java.time.Instant;
46+
import java.util.AbstractMap.SimpleEntry;
47+
import java.util.Collection;
48+
import java.util.Collections;
49+
import java.util.Date;
50+
import java.util.Map;
51+
import java.util.Map.Entry;
52+
import java.util.Set;
53+
import java.util.function.Supplier;
54+
import org.junit.jupiter.api.BeforeEach;
55+
import org.junit.jupiter.api.Test;
56+
import org.junit.jupiter.api.extension.ExtendWith;
57+
import org.mockito.ArgumentCaptor;
58+
import org.mockito.Captor;
59+
import org.mockito.Mock;
60+
import org.mockito.MockedStatic;
61+
import org.mockito.Mockito;
62+
import org.mockito.MockitoAnnotations;
63+
import org.mockito.junit.jupiter.MockitoExtension;
64+
65+
@ExtendWith(MockitoExtension.class)
66+
class GcpAuthAutoConfigurationCustomizerProviderTest {
67+
68+
@Mock private GoogleCredentials mockedGoogleCredentials;
69+
70+
@Captor private ArgumentCaptor<Collection<SpanData>> spanDataCollectionCaptor;
71+
72+
@Captor private ArgumentCaptor<Supplier<Map<String, String>>> headerSupplierCaptor;
73+
74+
private final Map<String, String> otelProperties =
75+
ImmutableMap.of(
76+
"otel.bsp.schedule.delay", // span exporter
77+
"10",
78+
"otel.traces.exporter",
79+
"otlp",
80+
"otel.metrics.exporter",
81+
"none",
82+
"otel.logs.exporter",
83+
"none");
84+
85+
@BeforeEach
86+
public void setup() {
87+
MockitoAnnotations.openMocks(this);
88+
Mockito.when(mockedGoogleCredentials.getQuotaProjectId()).thenReturn("test-project");
89+
Mockito.when(mockedGoogleCredentials.getAccessToken())
90+
.thenReturn(new AccessToken("fake", new Date()));
91+
}
92+
93+
@Test
94+
public void testCustomizerOtlpHttp() {
95+
OtlpHttpSpanExporter mockOtlpHttpSpanExporter = Mockito.mock(OtlpHttpSpanExporter.class);
96+
OtlpHttpSpanExporterBuilder otlpSpanExporterBuilder = OtlpHttpSpanExporter.builder();
97+
OtlpHttpSpanExporterBuilder spyOtlpHttpSpanExporterBuilder =
98+
Mockito.spy(otlpSpanExporterBuilder);
99+
Mockito.when(spyOtlpHttpSpanExporterBuilder.build()).thenReturn(mockOtlpHttpSpanExporter);
100+
101+
Mockito.when(mockOtlpHttpSpanExporter.export(Mockito.anyCollection()))
102+
.thenReturn(CompletableResultCode.ofSuccess());
103+
Mockito.when(mockOtlpHttpSpanExporter.toBuilder()).thenReturn(spyOtlpHttpSpanExporterBuilder);
104+
105+
try (MockedStatic<GoogleCredentials> googleCredentialsMockedStatic =
106+
Mockito.mockStatic(GoogleCredentials.class)) {
107+
googleCredentialsMockedStatic
108+
.when(GoogleCredentials::getApplicationDefault)
109+
.thenReturn(mockedGoogleCredentials);
110+
111+
OpenTelemetrySdk sdk = buildOpenTelemetrySdkWithExporter(mockOtlpHttpSpanExporter);
112+
generateTestSpan(sdk);
113+
114+
await()
115+
.atMost(Duration.ofSeconds(10))
116+
.untilAsserted(
117+
() -> {
118+
Mockito.verify(mockOtlpHttpSpanExporter, Mockito.times(1)).toBuilder();
119+
Mockito.verify(spyOtlpHttpSpanExporterBuilder, Mockito.times(1))
120+
.setHeaders(headerSupplierCaptor.capture());
121+
assertEquals(2, headerSupplierCaptor.getValue().get().size());
122+
verifyAuthHeaders(headerSupplierCaptor.getValue().get());
123+
124+
Mockito.verify(mockOtlpHttpSpanExporter, Mockito.atLeast(1))
125+
.export(spanDataCollectionCaptor.capture());
126+
spanDataCollectionCaptor
127+
.getValue()
128+
.forEach(
129+
spanData -> {
130+
assertEquals(
131+
"test-project",
132+
spanData
133+
.getAttributes()
134+
.get(AttributeKey.stringKey(GCP_USER_PROJECT_ID_KEY)));
135+
assertTrue(
136+
spanData
137+
.getAttributes()
138+
.asMap()
139+
.containsKey(AttributeKey.stringKey("work_loop")));
140+
});
141+
});
142+
}
143+
}
144+
145+
@Test
146+
public void testCustomizerOtlpGrpc() {
147+
OtlpGrpcSpanExporter mockOtlpGrpcSpanExporter = Mockito.mock(OtlpGrpcSpanExporter.class);
148+
OtlpGrpcSpanExporterBuilder otlpSpanExporterBuilder = OtlpGrpcSpanExporter.builder();
149+
OtlpGrpcSpanExporterBuilder spyOtlpGrpcSpanExporterBuilder =
150+
Mockito.spy(otlpSpanExporterBuilder);
151+
Mockito.when(spyOtlpGrpcSpanExporterBuilder.build()).thenReturn(mockOtlpGrpcSpanExporter);
152+
153+
Mockito.when(mockOtlpGrpcSpanExporter.export(Mockito.anyCollection()))
154+
.thenReturn(CompletableResultCode.ofSuccess());
155+
Mockito.when(mockOtlpGrpcSpanExporter.toBuilder()).thenReturn(spyOtlpGrpcSpanExporterBuilder);
156+
157+
try (MockedStatic<GoogleCredentials> googleCredentialsMockedStatic =
158+
Mockito.mockStatic(GoogleCredentials.class)) {
159+
googleCredentialsMockedStatic
160+
.when(GoogleCredentials::getApplicationDefault)
161+
.thenReturn(mockedGoogleCredentials);
162+
163+
OpenTelemetrySdk sdk = buildOpenTelemetrySdkWithExporter(mockOtlpGrpcSpanExporter);
164+
generateTestSpan(sdk);
165+
166+
await()
167+
.atMost(Duration.ofSeconds(10))
168+
.untilAsserted(
169+
() -> {
170+
Mockito.verify(mockOtlpGrpcSpanExporter, Mockito.times(1)).toBuilder();
171+
Mockito.verify(spyOtlpGrpcSpanExporterBuilder, Mockito.times(1))
172+
.setHeaders(headerSupplierCaptor.capture());
173+
assertEquals(2, headerSupplierCaptor.getValue().get().size());
174+
verifyAuthHeaders(headerSupplierCaptor.getValue().get());
175+
176+
Mockito.verify(mockOtlpGrpcSpanExporter, Mockito.atLeast(1))
177+
.export(spanDataCollectionCaptor.capture());
178+
spanDataCollectionCaptor
179+
.getValue()
180+
.forEach(
181+
spanData -> {
182+
assertEquals(
183+
"test-project",
184+
spanData
185+
.getAttributes()
186+
.get(AttributeKey.stringKey(GCP_USER_PROJECT_ID_KEY)));
187+
assertTrue(
188+
spanData
189+
.getAttributes()
190+
.asMap()
191+
.containsKey(AttributeKey.stringKey("work_loop")));
192+
});
193+
});
194+
}
195+
}
196+
197+
private OpenTelemetrySdk buildOpenTelemetrySdkWithExporter(SpanExporter spanExporter) {
198+
SpiHelper spiHelper =
199+
SpiHelper.create(GcpAuthAutoConfigurationCustomizerProviderTest.class.getClassLoader());
200+
AutoConfiguredOpenTelemetrySdkBuilder builder =
201+
AutoConfiguredOpenTelemetrySdk.builder().addPropertiesSupplier(() -> otelProperties);
202+
AutoConfigureUtil.setComponentLoader(
203+
builder,
204+
new ComponentLoader() {
205+
@SuppressWarnings("unchecked")
206+
@Override
207+
public <T> Iterable<T> load(Class<T> spiClass) {
208+
if (spiClass == ConfigurableSpanExporterProvider.class) {
209+
return Collections.singletonList(
210+
(T)
211+
new ConfigurableSpanExporterProvider() {
212+
@Override
213+
public SpanExporter createExporter(ConfigProperties configProperties) {
214+
return spanExporter;
215+
}
216+
217+
@Override
218+
public String getName() {
219+
return "otlp";
220+
}
221+
});
222+
}
223+
return spiHelper.load(spiClass);
224+
}
225+
});
226+
return builder.build().getOpenTelemetrySdk();
227+
}
228+
229+
private void verifyAuthHeaders(Map<String, String> headers) {
230+
Set<Entry<String, String>> headerEntrySet = headers.entrySet();
231+
assertTrue(
232+
headerEntrySet.contains(new SimpleEntry<>(QUOTA_USER_PROJECT_HEADER, "test-project")));
233+
assertTrue(headerEntrySet.contains(new SimpleEntry<>("Authorization", "Bearer fake")));
234+
}
235+
236+
private void generateTestSpan(OpenTelemetrySdk openTelemetrySdk) {
237+
Span span = openTelemetrySdk.getTracer("test").spanBuilder("sample").startSpan();
238+
try (Scope ignored = span.makeCurrent()) {
239+
long workOutput = busyloop();
240+
span.setAttribute("work_loop", workOutput);
241+
} finally {
242+
span.end();
243+
}
244+
}
245+
246+
// loop to simulate work done
247+
private long busyloop() {
248+
Instant start = Instant.now();
249+
Instant end;
250+
long counter = 0;
251+
do {
252+
counter++;
253+
end = Instant.now();
254+
} while (Duration.between(start, end).toMillis() < 1000);
255+
System.out.println("Busy work done, counted " + counter + " times in one second.");
256+
return counter;
257+
}
258+
}

0 commit comments

Comments
 (0)