chore: import upstream snapshot with attribution

tests/unittests/telemetry/test_metrics.py

@@ -0,0 +1,261 @@
+# 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.
+
+# pylint: disable=protected-access
+
+from unittest import mock
+
+from google.adk.telemetry import _metrics
+from google.genai import types
+from opentelemetry import metrics
+import pytest
+
+
+@pytest.fixture(name="mock_meter_setup")
+def _mock_meter_setup(monkeypatch):
+  """Sets up mock meter and histograms for testing."""
+  mock_meter = mock.MagicMock()
+  agent_duration_hist = mock.MagicMock(spec=metrics.Histogram)
+  workflow_duration_hist = mock.MagicMock(spec=metrics.Histogram)
+  tool_duration_hist = mock.MagicMock(spec=metrics.Histogram)
+  client_duration_hist = mock.MagicMock(spec=metrics.Histogram)
+  client_token_usage_hist = mock.MagicMock(spec=metrics.Histogram)
+
+  agent_duration_hist.name = "agent_invocation_duration"
+  workflow_duration_hist.name = "workflow_invocation_duration"
+  tool_duration_hist.name = "tool_execution_duration"
+  client_duration_hist.name = "client_operation_duration"
+  client_token_usage_hist.name = "client_token_usage"
+
+  def create_histogram_side_effect(name, **_kwargs):
+    if name == "gen_ai.invoke_agent.duration":
+      return agent_duration_hist
+    elif name == "gen_ai.invoke_workflow.duration":
+      return workflow_duration_hist
+    elif name == "gen_ai.execute_tool.duration":
+      return tool_duration_hist
+    elif name == "gen_ai.client.operation.duration":
+      return client_duration_hist
+    elif name == "gen_ai.client.token.usage":
+      return client_token_usage_hist
+    raise ValueError(f"Unknown metric name: {name}")
+
+  mock_meter.create_histogram.side_effect = create_histogram_side_effect
+
+  # Re-initialize the module-level variables in _metrics with mocked histograms
+  monkeypatch.setattr(_metrics, "meter", mock_meter)
+  monkeypatch.setattr(
+      _metrics, "_agent_invocation_duration", agent_duration_hist
+  )
+  monkeypatch.setattr(
+      _metrics, "_workflow_invocation_duration", workflow_duration_hist
+  )
+  monkeypatch.setattr(_metrics, "_tool_execution_duration", tool_duration_hist)
+  monkeypatch.setattr(
+      _metrics, "_client_operation_duration", client_duration_hist
+  )
+  monkeypatch.setattr(_metrics, "_client_token_usage", client_token_usage_hist)
+
+  return {
+      "meter": mock_meter,
+      "agent_duration": agent_duration_hist,
+      "workflow_duration": workflow_duration_hist,
+      "tool_duration": tool_duration_hist,
+      "client_duration": client_duration_hist,
+      "client_token_usage": client_token_usage_hist,
+  }
+
+
+def test_record_agent_invocation_duration(mock_meter_setup):
+  """Tests record_agent_invocation_duration records correctly."""
+  _metrics.record_agent_invocation_duration(
+      "test_agent",
+      1.0,
+  )
+  agent_duration_hist = mock_meter_setup["agent_duration"]
+  agent_duration_hist.record.assert_called_once()
+  args, kwargs = agent_duration_hist.record.call_args
+  assert args[0] == 1.0
+  want_attributes = {"gen_ai.agent.name": "test_agent"}
+  assert kwargs["attributes"] == want_attributes
+
+
+def test_record_agent_invocation_duration_with_error(mock_meter_setup):
+  """Tests record_agent_invocation_duration records error correctly."""
+  test_error = ValueError("agent failed")
+  _metrics.record_agent_invocation_duration(
+      "test_agent",
+      1.0,
+      error=test_error,
+  )
+  agent_duration_hist = mock_meter_setup["agent_duration"]
+  agent_duration_hist.record.assert_called_once()
+  _, kwargs = agent_duration_hist.record.call_args
+  assert kwargs["attributes"]["error.type"] == "ValueError"
+
+
+def test_record_workflow_invocation_duration_root(mock_meter_setup):
+  """Tests record_workflow_invocation_duration omits nested for the root."""
+  _metrics.record_workflow_invocation_duration(
+      workflow_name="my_workflow",
+      elapsed_s=1.0,
+      nested=False,
+  )
+  hist = mock_meter_setup["workflow_duration"]
+  hist.record.assert_called_once()
+  args, kwargs = hist.record.call_args
+  assert args[0] == 1.0
+  assert kwargs["attributes"] == {
+      "gen_ai.operation.name": "invoke_workflow",
+      "gen_ai.workflow.name": "my_workflow",
+  }
+
+
+def test_record_workflow_invocation_duration_nested_with_error(
+    mock_meter_setup,
+):
+  """Tests record_workflow_invocation_duration records nested + error."""
+  _metrics.record_workflow_invocation_duration(
+      workflow_name="nested_workflow",
+      elapsed_s=2.0,
+      nested=True,
+      error=ValueError("boom"),
+  )
+  hist = mock_meter_setup["workflow_duration"]
+  hist.record.assert_called_once()
+  _, kwargs = hist.record.call_args
+  assert kwargs["attributes"]["gen_ai.workflow.nested"] is True
+  assert kwargs["attributes"]["error.type"] == "ValueError"
+
+
+def test_record_tool_execution_duration(mock_meter_setup):
+  """Tests record_tool_execution_duration records correctly."""
+  _metrics.record_tool_execution_duration(
+      "test_tool",
+      "test_tool_type",
+      "test_agent",
+      0.5,
+  )
+  tool_duration_hist = mock_meter_setup["tool_duration"]
+  tool_duration_hist.record.assert_called_once()
+  args, kwargs = tool_duration_hist.record.call_args
+  assert args[0] == 0.5
+  want_attributes = {
+      "gen_ai.agent.name": "test_agent",
+      "gen_ai.tool.name": "test_tool",
+      "gen_ai.tool.type": "test_tool_type",
+  }
+  assert kwargs["attributes"] == want_attributes
+
+
+def test_record_tool_execution_duration_with_error(mock_meter_setup):
+  """Tests record_tool_execution_duration records error correctly."""
+  test_error = ValueError("tool failed")
+  _metrics.record_tool_execution_duration(
+      "test_tool",
+      "test_tool_type",
+      "test_agent",
+      0.5,
+      error=test_error,
+  )
+  tool_duration_hist = mock_meter_setup["tool_duration"]
+  tool_duration_hist.record.assert_called_once()
+  _, kwargs = tool_duration_hist.record.call_args
+  assert kwargs["attributes"]["error.type"] == "ValueError"
+
+
+def test_record_client_operation_duration(mock_meter_setup):
+  """Tests record_client_operation_duration records correctly."""
+  llm_request = mock.MagicMock(
+      contents=[types.Content(parts=[types.Part(text="hello")])]
+  )
+  response = mock.MagicMock(
+      content=types.Content(parts=[types.Part(text="hello response")])
+  )
+  _metrics.record_client_operation_duration(
+      agent_name="test_agent",
+      elapsed_s=0.1,
+      llm_request=llm_request,
+      responses=[response],
+  )
+  client_duration_hist = mock_meter_setup["client_duration"]
+  client_duration_hist.record.assert_called_once()
+  args, kwargs = client_duration_hist.record.call_args
+  assert args[0] == 0.1
+  want_attributes = {
+      "gen_ai.agent.name": "test_agent",
+      "gen_ai.operation.name": "generate_content",
+      "gen_ai.provider.name": "gemini",
+      "gen_ai.request.model": llm_request.model,
+      "gen_ai.response.model": response.model_version,
+  }
+  assert kwargs["attributes"] == want_attributes
+
+
+def test_record_client_token_usage(mock_meter_setup):
+  """Tests record_client_token_usage records correctly under different usage conditions."""
+  llm_request = mock.MagicMock(
+      contents=[types.Content(parts=[types.Part(text="hello")])],
+      model="test-model",
+  )
+  response = mock.MagicMock(
+      content=types.Content(parts=[types.Part(text="hello response")]),
+      model_version="test-model-v1",
+      usage_metadata=types.GenerateContentResponseUsageMetadata(
+          prompt_token_count=20,
+          candidates_token_count=30,
+          tool_use_prompt_token_count=5,
+          thoughts_token_count=10,
+      ),
+  )
+  _metrics.record_client_token_usage(
+      agent_name="test_agent",
+      llm_request=llm_request,
+      responses=[response],
+  )
+  client_token_usage_hist = mock_meter_setup["client_token_usage"]
+  assert client_token_usage_hist.record.call_count == 2
+
+  base_attributes = {
+      "gen_ai.agent.name": "test_agent",
+      "gen_ai.operation.name": "generate_content",
+      "gen_ai.provider.name": "gemini",
+      "gen_ai.request.model": "test-model",
+      "gen_ai.response.model": "test-model-v1",
+  }
+
+  input_call = None
+  output_call = None
+
+  for args, kwargs in client_token_usage_hist.record.call_args_list:
+    token_type = kwargs.get("attributes", {}).get("gen_ai.token.type")
+    if token_type == "input":
+      input_call = (args, kwargs)
+    elif token_type == "output":
+      output_call = (args, kwargs)
+
+  assert input_call is not None, "Missing 'input' token usage record"
+  assert output_call is not None, "Missing 'output' token usage record"
+
+  # Verify input tokens (prompt_token_count + tool_use_prompt_token_count)
+  assert input_call[0][0] == 25
+  assert input_call[1]["attributes"] == base_attributes | {
+      "gen_ai.token.type": "input"
+  }
+
+  # Verify output tokens (candidates_token_count + thoughts_token_count)
+  assert output_call[0][0] == 40
+  assert output_call[1]["attributes"] == base_attributes | {
+      "gen_ai.token.type": "output"
+  }