# Copyright 2026 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License 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. """Shared infrastructure for the telemetry functional tests. This module hosts: * The ``SpanDigest`` / ``LogDigest`` types used to build a deterministic comparison shape for in-memory spans + log records. * ``install_telemetry`` which patches an in-memory tracer + log exporter onto ADK's globals. * The canonical agent / tool / mock-LLM scenario shared across the ``test_functional.py``, ``test_node_functional.py`` and ``test_web_ui_functional.py`` test suites. * The ``FunctionalTestCase`` carrier used to parametrize tests against the hand-written expected shapes in ``functional_test_cases.py`` / ``functional_node_test_cases.py``. """ from __future__ import annotations from collections.abc import AsyncGenerator from collections.abc import Iterator from contextlib import aclosing from contextlib import contextmanager from dataclasses import dataclass from dataclasses import field from enum import Enum import gc import inspect import json import sys from types import CodeType from typing import Literal from typing import NamedTuple from typing import TYPE_CHECKING from google.adk.agents.llm_agent import Agent from google.adk.models.llm_response import LlmResponse from google.adk.runners import InMemoryRunner from google.adk.telemetry import _metrics from google.adk.telemetry import node_tracing from google.adk.telemetry import tracing from google.adk.tools.function_tool import FunctionTool from google.adk.workflow._base_node import START from google.adk.workflow._workflow import Workflow from google.genai.types import Content from google.genai.types import FinishReason from google.genai.types import Part from opentelemetry.sdk._logs import LoggerProvider from opentelemetry.sdk._logs.export import SimpleLogRecordProcessor from opentelemetry.sdk.metrics import MeterProvider from opentelemetry.sdk.metrics.export import HistogramDataPoint from opentelemetry.sdk.metrics.export import InMemoryMetricReader from opentelemetry.sdk.metrics.export import NumberDataPoint from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor import pytest if TYPE_CHECKING: from google.adk.events.event import Event from opentelemetry.sdk.trace import ReadableSpan from opentelemetry.util.types import AttributeValue from opentelemetry.sdk._logs import ReadableLogRecord from opentelemetry.sdk._logs.export import InMemoryLogRecordExporter from opentelemetry.sdk.metrics.export import MetricsData from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from ..testing_utils import MockModel from ..testing_utils import TestInMemoryRunner # --------------------------------------------------------------------------- # Env var + semconv constants. # --------------------------------------------------------------------------- OTEL_OPT_IN = "OTEL_SEMCONV_STABILITY_OPT_IN" CAPTURE_CONTENT = "OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT" EXPERIMENTAL_OPT_IN = "gen_ai_latest_experimental" ADK_TELEMETRY_SCHEMA_VERSION_OPT_IN = "ADK_TELEMETRY_SCHEMA_VERSION_OPT_IN" # Stable semconv event names. GEN_AI_SYSTEM_MESSAGE_EVENT = "gen_ai.system.message" GEN_AI_USER_MESSAGE_EVENT = "gen_ai.user.message" GEN_AI_CHOICE_EVENT = "gen_ai.choice" # Experimental semconv event name. GEN_AI_COMPLETION_DETAILS_EVENT = "gen_ai.client.inference.operation.details" # Difficult to extract, non deterministic attribute keys. # We check only for their presence, instead of their values. NON_DETERMINISTIC_ATTRIBUTE_KEYS: frozenset[str] = frozenset({ "gcp.vertex.agent.event_id", "gen_ai.tool.call.id", "gcp.vertex.agent.associated_event_ids", "gen_ai.conversation.id", "gcp.vertex.agent.invocation_id", "gcp.vertex.agent.session_id", }) # Span attribute keys whose values are JSON-serialized strings. # These are parsed back into Python objects before comparison so that JSON # property ordering doesn't drive test stability. JSON_ATTRIBUTE_KEYS: frozenset[str] = frozenset({ "gen_ai.input.messages", "gen_ai.output.messages", "gen_ai.system_instructions", "gen_ai.tool.definitions", }) # Sentinel used for non deterministic fields that we still want to assert as # being present. PRESENT = "PRESENT" # --------------------------------------------------------------------------- # Digests. # --------------------------------------------------------------------------- @dataclass(frozen=True) class LogDigest: """A deterministic digest of a ``ReadableLogRecord``. ``attributes`` and ``body`` are normalized via ``_normalize`` so test expectations can be written using plain Python literals (lists/dicts). """ event_name: str body: object = None attributes: dict[str, object] = field(default_factory=dict) @classmethod def from_log(cls, log: ReadableLogRecord) -> LogDigest: attrs: dict[str, object] = {} for k, v in (log.log_record.attributes or {}).items(): if k in NON_DETERMINISTIC_ATTRIBUTE_KEYS: attrs[k] = PRESENT else: attrs[k] = _normalize(v) return cls( event_name=log.log_record.event_name or "", body=_normalize(log.log_record.body), attributes=attrs, ) @dataclass(frozen=True) class SpanDigest: """A deterministic digest of a span in the in-memory span tree. In addition to the span's own name + attributes + child spans, each digest also carries the ``LogDigest`` records that were emitted while the span was the active span (matched by ``log_record.span_id``). """ name: str attributes: dict[str, AttributeValue] children: list[SpanDigest] = field(default_factory=list) logs: list[LogDigest] = field(default_factory=list) @classmethod def from_span(cls, span: ReadableSpan) -> SpanDigest: """Builds a single ``SpanDigest`` (no children, no logs) from a span. Attribute values are normalized so that: * Non-deterministic keys collapse to the ``PRESENT`` sentinel. * JSON-serialized attribute values are parsed into Python objects. * All other values pass through ``_normalize`` (tuples → lists, enums → ``.value``, ``None`` dict entries dropped). """ determinized_attributes: dict[str, AttributeValue] = {} for attr_key, attr_val in (span.attributes or {}).items(): if attr_key in NON_DETERMINISTIC_ATTRIBUTE_KEYS: determinized_attributes[attr_key] = PRESENT elif attr_key in JSON_ATTRIBUTE_KEYS and isinstance(attr_val, str): determinized_attributes[attr_key] = _normalize(json.loads(attr_val)) else: determinized_attributes[attr_key] = _normalize(attr_val) return cls(name=span.name, attributes=determinized_attributes) @classmethod def build( cls, spans: tuple[ReadableSpan, ...], logs: tuple[ReadableLogRecord, ...] = (), ) -> SpanDigest: """Builds the in-memory span tree, attaching logs by span id. Used for clear diffs with pytest assertions. """ digest_by_id: dict[int, SpanDigest] = {} for span in spans: if span.context is None: continue digest_by_id[span.context.span_id] = cls.from_span(span) # Attach each log to its enclosing span (matched by span_id). for log in logs: span_id = log.log_record.span_id if span_id is None or span_id == 0: continue digest = digest_by_id.get(span_id) if digest is None: continue digest.logs.append(LogDigest.from_log(log)) root: SpanDigest | None = None for span in spans: if span.context is None: continue digest = digest_by_id[span.context.span_id] if span.parent and span.parent.span_id in digest_by_id: parent_digest = digest_by_id[span.parent.span_id] parent_digest.children.append(digest) else: if root is not None: raise ValueError("Multiple root spans found.") root = digest # Sort for deterministic comparisons. for digest in digest_by_id.values(): digest.children.sort(key=lambda s: s.name) digest.logs[:] = sorted_log_digests(digest.logs) if root is None: raise ValueError("No root span found in the provided spans.") return root def all_logs(self) -> list[LogDigest]: """Returns all log digests in the tree, sorted deterministically.""" collected: list[LogDigest] = [] def _walk(node: SpanDigest) -> None: collected.extend(node.logs) for child in node.children: _walk(child) _walk(self) return sorted_log_digests(collected) def sorted_log_digests(logs: list[LogDigest]) -> list[LogDigest]: """Returns ``logs`` sorted in a stable, content-derived order.""" return sorted( logs, key=lambda log: ( log.event_name, json.dumps(log.body, sort_keys=True, default=str), json.dumps(log.attributes, sort_keys=True, default=str), ), ) class _NonDeterministic: """Sentinel for a metric value that is non-deterministic (e.g. wall-clock).""" __slots__ = () def __repr__(self) -> str: return "NON_DETERMINISTIC" # Marks a recorded metric value that cannot be pinned (e.g. ``*.duration`` # wall-clock timings); used in place of the actual value on both sides. NON_DETERMINISTIC = _NonDeterministic() @dataclass(frozen=True) class MetricPoint: """A single recorded metric data point.""" attributes: dict[str, AttributeValue] value: object def __hash__(self) -> int: return hash( (json.dumps(self.attributes, sort_keys=True, default=str), self.value) ) class HistogramSpec(NamedTuple): """Locates one ADK metric histogram so a test can redirect it. ``module`` is the module holding the histogram, ``attr`` the global on it to monkeypatch, and ``metric_name`` the instrument name it is recreated under. """ module: object attr: str metric_name: str # Histograms recorded by ADK. Each test redirects these onto an in-memory # reader so the recorded points can be asserted. _PATCHED_HISTOGRAMS: tuple[HistogramSpec, ...] = ( HistogramSpec( module=_metrics, attr="_agent_invocation_duration", metric_name="gen_ai.invoke_agent.duration", ), HistogramSpec( module=_metrics, attr="_tool_execution_duration", metric_name="gen_ai.execute_tool.duration", ), HistogramSpec( module=_metrics, attr="_client_operation_duration", metric_name="gen_ai.client.operation.duration", ), HistogramSpec( module=_metrics, attr="_client_token_usage", metric_name="gen_ai.client.token.usage", ), HistogramSpec( module=_metrics, attr="_workflow_invocation_duration", metric_name="gen_ai.invoke_workflow.duration", ), HistogramSpec( module=_metrics, attr="_invoke_agent_inference_calls", metric_name="gen_ai.invoke_agent.inference_calls", ), HistogramSpec( module=_metrics, attr="_invoke_agent_tool_calls", metric_name="gen_ai.invoke_agent.tool_calls", ), ) def _grouped_metric_points( metrics_data: MetricsData, ) -> dict[str, frozenset[MetricPoint]]: """Groups every recorded point by metric name as an order-free frozenset.""" grouped: dict[str, set[MetricPoint]] = {} for resource_metric in metrics_data.resource_metrics: for scope_metric in resource_metric.scope_metrics: for metric in scope_metric.metrics: for dp in metric.data.data_points: # Sum histograms expose ``.sum``; gauge / counter points expose # ``.value``. isinstance (not hasattr) keeps the typing precise. if isinstance(dp, HistogramDataPoint): value = dp.sum elif isinstance(dp, NumberDataPoint): value = dp.value else: value = NON_DETERMINISTIC # ``*.duration`` histograms record wall-clock timings, which are # non-deterministic; replace them so expectations need not pin a # timing. if metric.name.endswith(".duration"): value = NON_DETERMINISTIC grouped.setdefault(metric.name, set()).add( MetricPoint(attributes=dict(dp.attributes), value=value) ) return {name: frozenset(points) for name, points in grouped.items()} @dataclass(frozen=True) class TelemetryDigest: """The full telemetry surface produced by one scenario run. Bundles the root span tree (with per-span logs attached) and every recorded metric point grouped by metric name. Points are held in a frozenset per group so equality is independent of recording / authoring order. Test cases hand-write the expected instance; ``build`` produces the actual one. """ root_span: SpanDigest metric_points: dict[str, frozenset[MetricPoint]] @classmethod def build( cls, spans: tuple[ReadableSpan, ...], logs: tuple[ReadableLogRecord, ...], metrics_data: MetricsData, ) -> TelemetryDigest: """Builds the actual digest from in-memory spans, logs and metrics.""" return cls( root_span=SpanDigest.build(spans, logs), metric_points=_grouped_metric_points(metrics_data), ) def _normalize(value: object) -> object: """Normalizes a value for stable equality. * Tuples become lists (OTel coerces sequences to tuples on attributes). * Enums become their ``.value``. * Dict entries whose value is ``None`` are dropped (these are inserted by pydantic ``model_dump`` for unset fields and would dominate diffs). """ if isinstance(value, Enum): return value.value if isinstance(value, tuple): return [_normalize(v) for v in value] if isinstance(value, list): return [_normalize(v) for v in value] if isinstance(value, dict): return {k: _normalize(v) for k, v in value.items() if v is not None} return value # --------------------------------------------------------------------------- # Telemetry plumbing. # --------------------------------------------------------------------------- def install_telemetry( monkeypatch: pytest.MonkeyPatch, span_exporter: InMemorySpanExporter, log_exporter: InMemoryLogRecordExporter, metric_reader: InMemoryMetricReader, ) -> None: """Installs an in-memory tracer + log exporter + metric reader. Spans, logs and metric points emitted by ADK during the test are written into the provided exporters / reader. All three MUST be passed in so each test makes the choice of sink explicit (e.g. ``InMemoryLogRecordExporter`` vs ``WebUILogExporter``). """ tracer_provider = TracerProvider() tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter)) real_tracer = tracer_provider.get_tracer(__name__) monkeypatch.setattr( tracing.tracer, "start_as_current_span", real_tracer.start_as_current_span, ) monkeypatch.setattr( tracing.tracer, "start_span", real_tracer.start_span, ) monkeypatch.setattr( node_tracing.tracer, "start_as_current_span", real_tracer.start_as_current_span, ) monkeypatch.setattr( node_tracing.tracer, "start_span", real_tracer.start_span, ) logger_provider = LoggerProvider() logger_provider.add_log_record_processor( SimpleLogRecordProcessor(log_exporter) ) real_logger = logger_provider.get_logger(__name__) monkeypatch.setattr(tracing.otel_logger, "emit", real_logger.emit) meter_provider = MeterProvider(metric_readers=[metric_reader]) meter = meter_provider.get_meter("functional_test_meter") for spec in _PATCHED_HISTOGRAMS: monkeypatch.setattr( spec.module, spec.attr, meter.create_histogram(spec.metric_name) ) # --------------------------------------------------------------------------- # Canonical agent / tool / mock-LLM scenario. # --------------------------------------------------------------------------- USER_PROMPT = "hello" AGENT_NAME = "some_root_agent" AGENT_DESCRIPTION = "A sample root agent." BASE_INSTRUCTION = "you are helpful" # ADK auto-appends agent identity info to the system instruction when the # agent is invoked as the root of an InMemoryRunner directly. FULL_SYSTEM_INSTRUCTION = ( f"{BASE_INSTRUCTION}\n\n" f'You are an agent. Your internal name is "{AGENT_NAME}".' f' The description about you is "{AGENT_DESCRIPTION}".' ) FINAL_TEXT = "text response" TOOL_NAME = "some_tool" TOOL_DESCRIPTION = "A sample tool." TOOL_ARGS = {"arg1": "val1"} TOOL_RESULT_PREFIX = "processed " TOOL_RESULT = f"{TOOL_RESULT_PREFIX}{TOOL_ARGS['arg1']}" # The node scenario uses a workflow node whose output drives the agent's # input. The workflow itself wraps the same agent. WORKFLOW_NAME = "my_workflow" # The root workflow invokes a nested workflow whose sole node produces the # input for the agent. The nested workflow exercises the `gen_ai.workflow.nested` # span attribute + metric dimension (only nested workflows carry it). NESTED_WORKFLOW_NAME = "my_nested_workflow" NODE_NAME = "some_node" NODE_RESULT = "some result" NODE_USER_ID = "some_user" NODE_APP_NAME = "some_app" def _make_llm_response(part: Part) -> LlmResponse: return LlmResponse( content=Content(role="model", parts=[part]), finish_reason=FinishReason.STOP, ) def build_test_agent(*, failing: bool = False) -> Agent: """Builds the canonical 1-tool, 2-LLM-turn agent.""" mock_model = MockModel.create( responses=[ _make_llm_response( Part.from_function_call(name=TOOL_NAME, args=TOOL_ARGS) ), _make_llm_response(Part.from_text(text=FINAL_TEXT)), ] ) def some_tool(arg1: str) -> str: """A sample tool.""" if failing: raise ValueError("This tool always fails") return f"{TOOL_RESULT_PREFIX}{arg1}" return Agent( name=AGENT_NAME, description=AGENT_DESCRIPTION, instruction=BASE_INSTRUCTION, model=mock_model, tools=[FunctionTool(some_tool)], ) def build_test_runner(*, failing: bool = False) -> TestInMemoryRunner: """Builds a runner around the canonical agent (no workflow wrapper).""" return TestInMemoryRunner(node=build_test_agent(failing=failing)) def build_test_workflow(*, failing: bool = False) -> Workflow: """Builds the canonical Workflow: a nested workflow feeding the agent.""" test_agent = build_test_agent(failing=failing) async def some_node(ctx, node_input): return NODE_RESULT # Trivial workflow to test o11y of nested workflows nested_workflow = Workflow( name=NESTED_WORKFLOW_NAME, edges=[(START, some_node)], ) return Workflow( name=WORKFLOW_NAME, edges=[(START, nested_workflow, test_agent)], ) async def run_node_scenario( *, failing: bool = False, event_sink: list[Event] | None = None ) -> list[Event]: """Runs the workflow scenario to completion, draining the event stream. If ``event_sink`` is provided, collected events are appended to it as they are drained. This lets callers inspect the events that were emitted before an exception propagates (e.g. when ``failing=True``). """ workflow = build_test_workflow(failing=failing) runner = InMemoryRunner(app_name=NODE_APP_NAME, node=workflow) session = await runner.session_service.create_session( app_name=NODE_APP_NAME, user_id=NODE_USER_ID ) content = Content(parts=[Part.from_text(text=USER_PROMPT)], role="user") collected_events: list[Event] = event_sink if event_sink is not None else [] async with aclosing( runner.run_async( user_id=NODE_USER_ID, session_id=session.id, new_message=content, ) ) as agen: async for event in agen: collected_events.append(event) return collected_events async def run_agent_scenario(runner: TestInMemoryRunner) -> None: """Runs the non-node scenario to completion, draining the event stream.""" async with aclosing( runner.run_async_with_new_session_agen( Content(parts=[Part.from_text(text=USER_PROMPT)], role="user") ) ) as agen: async for _ in agen: pass # --------------------------------------------------------------------------- # Parametrization carrier. # --------------------------------------------------------------------------- @dataclass(frozen=True) class FunctionalTestCase: """One row of the (semconv, capture-content, schema-version) matrix.""" test_id: str semconv_opt_in: str | None capture_content: str | None schema_version: Literal[1, 2] expected: TelemetryDigest def apply_env(self, monkeypatch: pytest.MonkeyPatch) -> None: """Applies the per-case env vars for semconv + content capture. Always pins ``ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS=false`` so the tool span attributes remain deterministic across all cases. """ if self.semconv_opt_in is None: monkeypatch.delenv(OTEL_OPT_IN, raising=False) else: monkeypatch.setenv(OTEL_OPT_IN, self.semconv_opt_in) if self.capture_content is None: monkeypatch.delenv(CAPTURE_CONTENT, raising=False) else: monkeypatch.setenv(CAPTURE_CONTENT, self.capture_content) monkeypatch.setenv( ADK_TELEMETRY_SCHEMA_VERSION_OPT_IN, str(self.schema_version) ) monkeypatch.setenv("ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS", "false") # --------------------------------------------------------------------------- # aclosing wrapping assertions. # --------------------------------------------------------------------------- @contextmanager def aclosing_wrapping_assertions() -> Iterator[None]: """Context manager that asserts every async generator is wrapped in ``aclosing``. The check uses ``gc.get_referrers`` on every async generator first iterated within the block, which is expensive (~5 seconds per scenario). Run this once per scenario rather than per parametrized test case. On exit the original ``sys`` async-gen hooks are restored. """ prev_firstiter, prev_finalizer = sys.get_asyncgen_hooks() def wrapped_firstiter(coro: AsyncGenerator[object, object]): if _is_async_context_manager(): if prev_firstiter: prev_firstiter(coro) return assert any( isinstance(referrer, aclosing) or isinstance(indirect_referrer, aclosing) for referrer in gc.get_referrers(coro) # Some coroutines have a layer of indirection in Python 3.10 for indirect_referrer in gc.get_referrers(referrer) ), _no_aclosing_assertion_error(coro) if prev_firstiter: prev_firstiter(coro) sys.set_asyncgen_hooks(wrapped_firstiter, prev_finalizer) try: yield finally: sys.set_asyncgen_hooks(prev_firstiter, prev_finalizer) def _no_aclosing_assertion_error(coro: AsyncGenerator[object, object]) -> str: first_iter_loc = "" definition_loc = "" if (f := inspect.currentframe()) and (f := f.f_back) and (f := f.f_back): first_iter_loc = f'file "{f.f_code.co_filename}" line "{f.f_lineno}"' if (ag_code := getattr(coro, "ag_code", None)) and isinstance( ag_code, CodeType ): definition_loc = ( f'file "{ag_code.co_filename}" line "{ag_code.co_firstlineno}"' ) header_str = f'Async generator "{coro.__name__}" is not wrapped in aclosing' first_iter_str = ( f"first iterated in {first_iter_loc}" if first_iter_loc else "" ) definition_str = f"defined in {definition_loc}" if definition_loc else "" instruction_str = """ Wrap the iteration in the following code snippet before iterating: async with contextlib.aclosing(...) as agen: async for ... as agen: ... """ return "\n".join( part for part in [ header_str, first_iter_str, definition_str, instruction_str, ] if part ) def _is_async_context_manager() -> bool: """Checks if this function was invoked by contextlib.asynccontextmanager.""" frame = inspect.currentframe() while frame: if ( frame.f_code.co_name == "__aenter__" and "contextlib" in frame.f_code.co_filename ): return True frame = frame.f_back return False