chore: import upstream snapshot with attribution

tests/unittests/tools/data_agent/test_data_agent_tool.py

@@ -0,0 +1,224 @@
+# 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.
+
+from unittest import mock
+
+from google.adk.tools.data_agent import data_agent_tool
+from google.adk.tools.tool_context import ToolContext
+
+
+@mock.patch.object(
+    data_agent_tool._gda_stream_util, "get_gda_session", autospec=True
+)
+def test_list_accessible_data_agents_success(mock_get_session):
+  """Tests list_accessible_data_agents success path."""
+  mock_creds = mock.Mock()
+  mock_session = mock.MagicMock()
+  mock_response = mock.Mock()
+  mock_response.json.return_value = {"dataAgents": ["agent1", "agent2"]}
+  mock_response.raise_for_status.return_value = None
+  mock_session.get.return_value = mock_response
+  mock_get_session.return_value = (
+      mock_session,
+      "https://geminidataanalytics.googleapis.com",
+  )
+  result = data_agent_tool.list_accessible_data_agents(
+      "test-project", mock_creds
+  )
+  assert result["status"] == "SUCCESS"
+  assert result["response"] == ["agent1", "agent2"]
+  mock_get_session.assert_called_once_with(mock_creds)
+  mock_session.get.assert_called_once_with(
+      "https://geminidataanalytics.googleapis.com/v1beta/projects/test-project/locations/global/dataAgents:listAccessible",
+      headers={
+          "Content-Type": "application/json",
+          "X-Goog-API-Client": "GOOGLE_ADK",
+      },
+  )
+
+
+@mock.patch.object(
+    data_agent_tool._gda_stream_util, "get_gda_session", autospec=True
+)
+def test_list_accessible_data_agents_exception(mock_get_session):
+  """Tests list_accessible_data_agents exception path."""
+  mock_creds = mock.Mock()
+  mock_session = mock.MagicMock()
+  mock_session.get.side_effect = Exception("List failed!")
+  mock_get_session.return_value = (
+      mock_session,
+      "https://geminidataanalytics.googleapis.com",
+  )
+  result = data_agent_tool.list_accessible_data_agents(
+      "test-project", mock_creds
+  )
+  assert result["status"] == "ERROR"
+  assert "List failed!" in result["error_details"]
+  mock_get_session.assert_called_once_with(mock_creds)
+  mock_session.get.assert_called_once()
+
+
+@mock.patch.object(
+    data_agent_tool._gda_stream_util, "get_gda_endpoint", autospec=True
+)
+@mock.patch.object(
+    data_agent_tool._gda_stream_util, "get_gda_session", autospec=True
+)
+def test_get_data_agent_info_success(mock_get_session, mock_get_endpoint):
+  """Tests get_data_agent_info success path."""
+  mock_creds = mock.Mock()
+  mock_session = mock.MagicMock()
+  mock_response = mock.Mock()
+  mock_response.json.return_value = "agent_info"
+  mock_response.raise_for_status.return_value = None
+  mock_session.get.return_value = mock_response
+  mock_get_endpoint.return_value = "https://geminidataanalytics.googleapis.com"
+  mock_get_session.return_value = (
+      mock_session,
+      "https://geminidataanalytics.googleapis.com",
+  )
+  result = data_agent_tool.get_data_agent_info("agent_name", mock_creds)
+  assert result["status"] == "SUCCESS"
+  assert result["response"] == "agent_info"
+  mock_get_session.assert_called_once_with(mock_creds)
+  mock_get_endpoint.assert_called_once()
+  mock_session.get.assert_called_once_with(
+      "https://geminidataanalytics.googleapis.com/v1beta/agent_name",
+      headers={
+          "Content-Type": "application/json",
+          "X-Goog-API-Client": "GOOGLE_ADK",
+      },
+  )
+
+
+@mock.patch.object(
+    data_agent_tool._gda_stream_util, "get_gda_endpoint", autospec=True
+)
+@mock.patch.object(
+    data_agent_tool._gda_stream_util, "get_gda_session", autospec=True
+)
+def test_get_data_agent_info_exception(mock_get_session, mock_get_endpoint):
+  """Tests get_data_agent_info exception path."""
+  mock_creds = mock.Mock()
+  mock_session = mock.MagicMock()
+  mock_session.get.side_effect = Exception("Get failed!")
+  mock_get_endpoint.return_value = "https://geminidataanalytics.googleapis.com"
+  mock_get_session.return_value = (
+      mock_session,
+      "https://geminidataanalytics.googleapis.com",
+  )
+  result = data_agent_tool.get_data_agent_info("agent_name", mock_creds)
+  assert result["status"] == "ERROR"
+  assert "Get failed!" in result["error_details"]
+  mock_get_session.assert_called_once_with(mock_creds)
+  mock_get_endpoint.assert_called_once()
+  mock_session.get.assert_called_once()
+
+
+@mock.patch.object(
+    data_agent_tool._gda_stream_util, "get_stream", autospec=True
+)
+@mock.patch.object(
+    data_agent_tool._gda_stream_util, "get_gda_session", autospec=True
+)
+@mock.patch.object(data_agent_tool, "_get_data_agent_info", autospec=True)
+def test_ask_data_agent_success(
+    mock_get_agent_info, mock_get_session, mock_get_stream
+):
+  """Tests ask_data_agent success path."""
+  mock_creds = mock.Mock()
+  mock_session = mock.MagicMock()
+  mock_get_session.return_value = (
+      mock_session,
+      "https://geminidataanalytics.googleapis.com",
+  )
+  mock_get_agent_info.return_value = {"status": "SUCCESS", "response": {}}
+  mock_get_stream.return_value = [
+      {"text": {"parts": ["response1"], "textType": "THOUGHT"}},
+      {"text": {"parts": ["response2"], "textType": "FINAL_RESPONSE"}},
+  ]
+  mock_invocation_context = mock.Mock()
+  mock_invocation_context.session.state = {}
+  mock_context = ToolContext(mock_invocation_context)
+  mock_settings = mock.Mock()
+
+  result = data_agent_tool.ask_data_agent(
+      "projects/p/locations/l/dataAgents/a",
+      "query",
+      credentials=mock_creds,
+      tool_context=mock_context,
+      settings=mock_settings,
+  )
+  assert result["status"] == "SUCCESS"
+  assert result["response"] == [
+      {"text": {"parts": ["response1"], "textType": "THOUGHT"}},
+      {"text": {"parts": ["response2"], "textType": "FINAL_RESPONSE"}},
+  ]
+  mock_get_agent_info.assert_called_once_with(
+      "projects/p/locations/l/dataAgents/a", mock_creds, session=mock_session
+  )
+  mock_get_session.assert_called_once_with(mock_creds)
+  mock_get_stream.assert_called_once_with(
+      mock_session,
+      "https://geminidataanalytics.googleapis.com/v1beta/projects/p/locations/l:chat",
+      {
+          "messages": [{"userMessage": {"text": "query"}}],
+          "dataAgentContext": {
+              "dataAgent": "projects/p/locations/l/dataAgents/a",
+          },
+          "clientIdEnum": "GOOGLE_ADK",
+      },
+      {
+          "Content-Type": "application/json",
+          "X-Goog-API-Client": "GOOGLE_ADK",
+      },
+      mock_settings.max_query_result_rows,
+  )
+
+
+@mock.patch.object(
+    data_agent_tool._gda_stream_util, "get_stream", autospec=True
+)
+@mock.patch.object(
+    data_agent_tool._gda_stream_util, "get_gda_session", autospec=True
+)
+@mock.patch.object(data_agent_tool, "_get_data_agent_info", autospec=True)
+def test_ask_data_agent_exception(
+    mock_get_agent_info, mock_get_session, mock_get_stream
+):
+  """Tests ask_data_agent exception path."""
+  mock_creds = mock.Mock()
+  mock_session = mock.MagicMock()
+  mock_get_session.return_value = (
+      mock_session,
+      "https://geminidataanalytics.googleapis.com",
+  )
+  mock_get_agent_info.return_value = {"status": "SUCCESS", "response": {}}
+  mock_get_stream.side_effect = Exception("Chat failed!")
+  mock_invocation_context = mock.Mock()
+  mock_invocation_context.session.state = {}
+  mock_context = ToolContext(mock_invocation_context)
+  mock_settings = mock.Mock()
+
+  result = data_agent_tool.ask_data_agent(
+      "projects/p/locations/l/dataAgents/a",
+      "query",
+      credentials=mock_creds,
+      tool_context=mock_context,
+      settings=mock_settings,
+  )
+  assert result["status"] == "ERROR"
+  assert "Chat failed!" in result["error_details"]
+  mock_get_session.assert_called_once_with(mock_creds)
+  mock_get_stream.assert_called_once()

tests/unittests/tools/data_agent/test_data_agent_toolset.py

@@ -0,0 +1,128 @@
+# 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.
+
+from __future__ import annotations
+
+from unittest import mock
+
+from google.adk.tools.data_agent import DataAgentCredentialsConfig
+from google.adk.tools.data_agent import DataAgentToolset
+from google.adk.tools.data_agent.config import DataAgentToolConfig
+from google.adk.tools.google_tool import GoogleTool
+import pytest
+
+
+@pytest.mark.asyncio
+async def test_data_agent_toolset_tools_default():
+  """Test default DataAgentToolset.
+
+  This test verifies the behavior of the DataAgentToolset when no filter is
+  specified.
+  """
+  credentials_config = DataAgentCredentialsConfig(
+      client_id="abc", client_secret="def"
+  )
+  toolset = DataAgentToolset(
+      credentials_config=credentials_config, data_agent_tool_config=None
+  )
+  # Verify that the tool config is initialized to default values.
+  assert isinstance(toolset._tool_settings, DataAgentToolConfig)  # pylint: disable=protected-access
+  assert toolset._tool_settings.__dict__ == DataAgentToolConfig().__dict__  # pylint: disable=protected-access
+
+  tools = await toolset.get_tools()
+  assert tools is not None
+
+  assert len(tools) == 3
+  assert all(isinstance(tool, GoogleTool) for tool in tools)
+
+  expected_tool_names = set([
+      "list_accessible_data_agents",
+      "get_data_agent_info",
+      "ask_data_agent",
+  ])
+  actual_tool_names = {tool.name for tool in tools}
+  assert actual_tool_names == expected_tool_names
+
+
+@pytest.mark.parametrize(
+    "selected_tools",
+    [
+        pytest.param([], id="None"),
+        pytest.param(
+            ["list_accessible_data_agents", "get_data_agent_info"],
+            id="list_and_get",
+        ),
+        pytest.param(["ask_data_agent"], id="ask"),
+    ],
+)
+@pytest.mark.asyncio
+async def test_data_agent_toolset_tools_selective(selected_tools):
+  """Test DataAgentToolset with filter.
+
+  This test verifies the behavior of the DataAgentToolset when filter is
+  specified. A use case for this would be when the agent builder wants to
+  use only a subset of the tools provided by the toolset.
+  """
+  credentials_config = DataAgentCredentialsConfig(
+      client_id="abc", client_secret="def"
+  )
+  toolset = DataAgentToolset(
+      credentials_config=credentials_config, tool_filter=selected_tools
+  )
+  tools = await toolset.get_tools()
+  assert tools is not None
+
+  assert len(tools) == len(selected_tools)
+  assert all(isinstance(tool, GoogleTool) for tool in tools)
+
+  expected_tool_names = set(selected_tools)
+  actual_tool_names = {tool.name for tool in tools}
+  assert actual_tool_names == expected_tool_names
+
+
+@pytest.mark.parametrize(
+    ("selected_tools", "returned_tools"),
+    [
+        pytest.param(["unknown"], [], id="all-unknown"),
+        pytest.param(
+            ["unknown", "ask_data_agent"],
+            ["ask_data_agent"],
+            id="mixed-known-unknown",
+        ),
+    ],
+)
+@pytest.mark.asyncio
+async def test_data_agent_toolset_unknown_tool(selected_tools, returned_tools):
+  """Test DataAgentToolset with filter.
+
+  This test verifies the behavior of the DataAgentToolset when filter is
+  specified with an unknown tool.
+  """
+  credentials_config = DataAgentCredentialsConfig(
+      client_id="abc", client_secret="def"
+  )
+
+  toolset = DataAgentToolset(
+      credentials_config=credentials_config, tool_filter=selected_tools
+  )
+
+  tools = await toolset.get_tools()
+  assert tools is not None
+
+  assert len(tools) == len(returned_tools)
+  assert all(isinstance(tool, GoogleTool) for tool in tools)
+
+  expected_tool_names = set(returned_tools)
+  actual_tool_names = {tool.name for tool in tools}
+  assert actual_tool_names == expected_tool_names