test: Adding a test for telemetry in RunnerTest

google-genai-bot · copybara-github · commit 77ebde8c4e5a · 2025-10-22T10:29:37.000-07:00
PiperOrigin-RevId: 822575688
diff --git a/core/pom.xml b/core/pom.xml
@@ -184,6 +184,10 @@
       <groupId>io.opentelemetry</groupId>
       <artifactId>opentelemetry-sdk-trace</artifactId>
     </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk-testing</artifactId>
+    </dependency>
   </dependencies>
   <build>
     <resources>
diff --git a/core/src/main/java/com/google/adk/Telemetry.java b/core/src/main/java/com/google/adk/Telemetry.java
@@ -46,10 +46,17 @@
 public class Telemetry {
 
   private static final Logger log = LoggerFactory.getLogger(Telemetry.class);
-  private static final Tracer tracer = GlobalOpenTelemetry.getTracer("gcp.vertex.agent");
+
+  @SuppressWarnings("NonFinalStaticField")
+  private static Tracer tracer = GlobalOpenTelemetry.getTracer("gcp.vertex.agent");
 
   private Telemetry() {}
 
+  /** Sets the OpenTelemetry instance to be used for tracing. This is for testing purposes only. */
+  public static void setTracerForTesting(Tracer tracer) {
+    Telemetry.tracer = tracer;
+  }
+
   /**
    * Traces tool call arguments.
    *
diff --git a/core/src/test/java/com/google/adk/runner/RunnerTest.java b/core/src/test/java/com/google/adk/runner/RunnerTest.java
@@ -28,7 +28,9 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import com.google.adk.Telemetry;
 import com.google.adk.agents.InvocationContext;
+import com.google.adk.agents.LiveRequestQueue;
 import com.google.adk.agents.LlmAgent;
 import com.google.adk.agents.RunConfig;
 import com.google.adk.events.Event;
@@ -44,19 +46,27 @@
 import com.google.genai.types.Content;
 import com.google.genai.types.FunctionCall;
 import com.google.genai.types.Part;
+import io.opentelemetry.api.trace.Tracer;
+import io.opentelemetry.sdk.testing.junit4.OpenTelemetryRule;
+import io.opentelemetry.sdk.trace.data.SpanData;
 import io.reactivex.rxjava3.core.Completable;
 import io.reactivex.rxjava3.core.Flowable;
 import io.reactivex.rxjava3.core.Maybe;
 import java.util.List;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 import org.mockito.ArgumentCaptor;
 
 @RunWith(JUnit4.class)
 public final class RunnerTest {
+  @Rule public final OpenTelemetryRule openTelemetryRule = OpenTelemetryRule.create();
 
   private final BasePlugin plugin = mockPlugin("test");
   private final Content pluginContent = createContent("from plugin");
@@ -65,6 +75,7 @@ public final class RunnerTest {
   private final Runner runner = new InMemoryRunner(agent, "test", ImmutableList.of(plugin));
   private final Session session =
       runner.sessionService().createSession("test", "user").blockingGet();
+  private Tracer originalTracer;
 
   private final FailingEchoTool failingEchoTool = new FailingEchoTool();
   private final EchoTool echoTool = new EchoTool();
@@ -94,6 +105,17 @@ private BasePlugin mockPlugin(String name) {
     return plugin;
   }
 
+  @Before
+  public void setUp() {
+    this.originalTracer = Telemetry.getTracer();
+    Telemetry.setTracerForTesting(openTelemetryRule.getOpenTelemetry().getTracer("RunnerTest"));
+  }
+
+  @After
+  public void tearDown() {
+    Telemetry.setTracerForTesting(originalTracer);
+  }
+
   @Test
   public void pluginDoesNothing() {
     var events =
@@ -606,4 +628,34 @@ public void beforeRunCallback_withStateDelta_seesMergedState() {
   private Content createContent(String text) {
     return Content.builder().parts(Part.builder().text(text).build()).build();
   }
+
+  @Test
+  public void runAsync_createsInvocationSpan() {
+    var unused =
+        runner.runAsync("user", session.id(), createContent("test message")).toList().blockingGet();
+
+    List<SpanData> spans = openTelemetryRule.getSpans();
+    assertThat(spans).isNotEmpty();
+
+    Optional<SpanData> invocationSpan =
+        spans.stream().filter(span -> Objects.equals(span.getName(), "invocation")).findFirst();
+
+    assertThat(invocationSpan).isPresent();
+    assertThat(invocationSpan.get().hasEnded()).isTrue();
+  }
+
+  @Test
+  public void runLive_createsInvocationSpan() {
+    LiveRequestQueue liveRequestQueue = new LiveRequestQueue();
+    var unused = runner.runLive(session, liveRequestQueue, RunConfig.builder().build()).test();
+
+    List<SpanData> spans = openTelemetryRule.getSpans();
+    assertThat(spans).isNotEmpty();
+
+    Optional<SpanData> invocationSpan =
+        spans.stream().filter(span -> Objects.equals(span.getName(), "invocation")).findFirst();
+
+    assertThat(invocationSpan).isPresent();
+    assertThat(invocationSpan.get().hasEnded()).isTrue();
+  }
 }
