google--adk-python/tests/unittests/tools/spanner/test_search_tool.py

# 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 unittest.mock import MagicMock

from google.adk.tools.spanner import client
from google.adk.tools.spanner import search_tool
from google.adk.tools.spanner import utils
from google.cloud.spanner_admin_database_v1.types import DatabaseDialect
import pytest


@pytest.fixture
def mock_credentials():
  return MagicMock()


@pytest.fixture
def mock_spanner_ids():
  return {
      "project_id": "test-project",
      "instance_id": "test-instance",
      "database_id": "test-database",
      "table_name": "test-table",
  }


@pytest.mark.parametrize(
    ("embedding_option_key", "embedding_option_value", "expected_embedding"),
    [
        pytest.param(
            "spanner_googlesql_embedding_model_name",
            "EmbeddingsModel",
            [0.1, 0.2, 0.3],
            id="spanner_googlesql_embedding_model",
        ),
        pytest.param(
            "vertex_ai_embedding_model_name",
            "text-embedding-005",
            [0.4, 0.5, 0.6],
            id="vertex_ai_embedding_model",
        ),
    ],
)
@pytest.mark.asyncio
@mock.patch.object(utils, "embed_contents_async", autospec=True)
@mock.patch.object(client, "get_spanner_client")
async def test_similarity_search_knn_success(
    mock_get_spanner_client,
    mock_embed_contents_async,
    mock_spanner_ids,
    mock_credentials,
    embedding_option_key,
    embedding_option_value,
    expected_embedding,
):
  """Test similarity_search function with kNN success."""
  mock_spanner_client = MagicMock()
  mock_instance = MagicMock()
  mock_database = MagicMock()
  mock_snapshot = MagicMock()
  mock_database.snapshot.return_value.__enter__.return_value = mock_snapshot
  mock_database.database_dialect = DatabaseDialect.GOOGLE_STANDARD_SQL
  mock_instance.database.return_value = mock_database
  mock_spanner_client.instance.return_value = mock_instance
  mock_get_spanner_client.return_value = mock_spanner_client

  if embedding_option_key == "vertex_ai_embedding_model_name":
    mock_embed_contents_async.return_value = [expected_embedding]
    # execute_sql is called once for the kNN search
    mock_snapshot.execute_sql.return_value = iter([("result1",), ("result2",)])
  else:
    mock_embedding_result = MagicMock()
    mock_embedding_result.one.return_value = (expected_embedding,)
    # First call to execute_sql is for getting the embedding,
    # second call is for the kNN search
    mock_snapshot.execute_sql.side_effect = [
        mock_embedding_result,
        iter([("result1",), ("result2",)]),
    ]

  result = await search_tool.similarity_search(
      project_id=mock_spanner_ids["project_id"],
      instance_id=mock_spanner_ids["instance_id"],
      database_id=mock_spanner_ids["database_id"],
      table_name=mock_spanner_ids["table_name"],
      query="test query",
      embedding_column_to_search="embedding_col",
      columns=["col1"],
      embedding_options={embedding_option_key: embedding_option_value},
      credentials=mock_credentials,
  )
  assert result["status"] == "SUCCESS", result
  assert result["rows"] == [("result1",), ("result2",)]

  # Check the generated SQL for kNN search
  call_args = mock_snapshot.execute_sql.call_args
  sql = call_args.args[0]
  assert "COSINE_DISTANCE" in sql
  assert "@embedding" in sql
  assert call_args.kwargs == {"params": {"embedding": expected_embedding}}
  if embedding_option_key == "vertex_ai_embedding_model_name":
    mock_embed_contents_async.assert_called_once_with(
        embedding_option_value, ["test query"], None
    )


@pytest.mark.asyncio
@mock.patch.object(client, "get_spanner_client")
async def test_similarity_search_ann_success(
    mock_get_spanner_client, mock_spanner_ids, mock_credentials
):
  """Test similarity_search function with ANN success."""
  mock_spanner_client = MagicMock()
  mock_instance = MagicMock()
  mock_database = MagicMock()
  mock_snapshot = MagicMock()
  mock_embedding_result = MagicMock()
  mock_embedding_result.one.return_value = ([0.1, 0.2, 0.3],)
  # First call to execute_sql is for getting the embedding
  # Second call is for the ANN search
  mock_snapshot.execute_sql.side_effect = [
      mock_embedding_result,
      iter([("ann_result1",), ("ann_result2",)]),
  ]
  mock_database.snapshot.return_value.__enter__.return_value = mock_snapshot
  mock_database.database_dialect = DatabaseDialect.GOOGLE_STANDARD_SQL
  mock_instance.database.return_value = mock_database
  mock_spanner_client.instance.return_value = mock_instance
  mock_get_spanner_client.return_value = mock_spanner_client

  result = await search_tool.similarity_search(
      project_id=mock_spanner_ids["project_id"],
      instance_id=mock_spanner_ids["instance_id"],
      database_id=mock_spanner_ids["database_id"],
      table_name=mock_spanner_ids["table_name"],
      query="test query",
      embedding_column_to_search="embedding_col",
      columns=["col1"],
      embedding_options={
          "spanner_googlesql_embedding_model_name": "test_model"
      },
      credentials=mock_credentials,
      search_options={
          "nearest_neighbors_algorithm": "APPROXIMATE_NEAREST_NEIGHBORS"
      },
  )
  assert result["status"] == "SUCCESS", result
  assert result["rows"] == [("ann_result1",), ("ann_result2",)]
  call_args = mock_snapshot.execute_sql.call_args
  sql = call_args.args[0]
  assert "APPROX_COSINE_DISTANCE" in sql
  assert "@embedding" in sql
  assert call_args.kwargs == {"params": {"embedding": [0.1, 0.2, 0.3]}}


@pytest.mark.asyncio
@mock.patch.object(client, "get_spanner_client")
async def test_similarity_search_error(
    mock_get_spanner_client, mock_spanner_ids, mock_credentials
):
  """Test similarity_search function with a generic error."""
  mock_get_spanner_client.side_effect = Exception("Test Exception")
  result = await search_tool.similarity_search(
      project_id=mock_spanner_ids["project_id"],
      instance_id=mock_spanner_ids["instance_id"],
      database_id=mock_spanner_ids["database_id"],
      table_name=mock_spanner_ids["table_name"],
      query="test query",
      embedding_column_to_search="embedding_col",
      embedding_options={
          "spanner_googlesql_embedding_model_name": "test_model"
      },
      columns=["col1"],
      credentials=mock_credentials,
  )
  assert result["status"] == "ERROR"
  assert "Test Exception" in result["error_details"]


@pytest.mark.asyncio
@mock.patch.object(utils, "embed_contents_async")
@mock.patch.object(client, "get_spanner_client")
async def test_similarity_search_circular_row_fallback_to_string(
    mock_get_spanner_client,
    mock_embed_contents_async,
    mock_spanner_ids,
    mock_credentials,
):
  """Test similarity_search stringifies rows with circular references."""
  mock_spanner_client = MagicMock()
  mock_instance = MagicMock()
  mock_database = MagicMock()
  mock_snapshot = MagicMock()
  circular_row = []
  circular_row.append(circular_row)
  mock_embed_contents_async.return_value = [[0.1, 0.2, 0.3]]
  mock_snapshot.execute_sql.return_value = iter([circular_row])
  mock_database.snapshot.return_value.__enter__.return_value = mock_snapshot
  mock_database.database_dialect = DatabaseDialect.GOOGLE_STANDARD_SQL
  mock_instance.database.return_value = mock_database
  mock_spanner_client.instance.return_value = mock_instance
  mock_get_spanner_client.return_value = mock_spanner_client

  result = await search_tool.similarity_search(
      project_id=mock_spanner_ids["project_id"],
      instance_id=mock_spanner_ids["instance_id"],
      database_id=mock_spanner_ids["database_id"],
      table_name=mock_spanner_ids["table_name"],
      query="test query",
      embedding_column_to_search="embedding_col",
      columns=["col1"],
      embedding_options={
          "vertex_ai_embedding_model_name": "text-embedding-005"
      },
      credentials=mock_credentials,
  )

  assert result["status"] == "SUCCESS", result
  assert result["rows"] == [str(circular_row)]


@pytest.mark.asyncio
@mock.patch.object(client, "get_spanner_client")
async def test_similarity_search_postgresql_knn_success(
    mock_get_spanner_client, mock_spanner_ids, mock_credentials
):
  """Test similarity_search with PostgreSQL dialect for kNN."""
  mock_spanner_client = MagicMock()
  mock_instance = MagicMock()
  mock_database = MagicMock()
  mock_snapshot = MagicMock()
  mock_embedding_result = MagicMock()
  mock_embedding_result.one.return_value = ([0.1, 0.2, 0.3],)
  mock_snapshot.execute_sql.side_effect = [
      mock_embedding_result,
      iter([("pg_result",)]),
  ]
  mock_database.snapshot.return_value.__enter__.return_value = mock_snapshot
  mock_database.database_dialect = DatabaseDialect.POSTGRESQL
  mock_instance.database.return_value = mock_database
  mock_spanner_client.instance.return_value = mock_instance
  mock_get_spanner_client.return_value = mock_spanner_client

  result = await search_tool.similarity_search(
      project_id=mock_spanner_ids["project_id"],
      instance_id=mock_spanner_ids["instance_id"],
      database_id=mock_spanner_ids["database_id"],
      table_name=mock_spanner_ids["table_name"],
      query="test query",
      embedding_column_to_search="embedding_col",
      columns=["col1"],
      embedding_options={
          "spanner_postgresql_vertex_ai_embedding_model_endpoint": (
              "test_endpoint"
          )
      },
      credentials=mock_credentials,
  )
  assert result["status"] == "SUCCESS", result
  assert result["rows"] == [("pg_result",)]
  call_args = mock_snapshot.execute_sql.call_args
  sql = call_args.args[0]
  assert "spanner.cosine_distance" in sql
  assert "$1" in sql
  assert call_args.kwargs == {"params": {"p1": [0.1, 0.2, 0.3]}}


@pytest.mark.asyncio
@mock.patch.object(client, "get_spanner_client")
async def test_similarity_search_postgresql_ann_unsupported(
    mock_get_spanner_client, mock_spanner_ids, mock_credentials
):
  """Test similarity_search with unsupported ANN for PostgreSQL dialect."""
  mock_spanner_client = MagicMock()
  mock_instance = MagicMock()
  mock_database = MagicMock()
  mock_database.database_dialect = DatabaseDialect.POSTGRESQL
  mock_instance.database.return_value = mock_database
  mock_spanner_client.instance.return_value = mock_instance
  mock_get_spanner_client.return_value = mock_spanner_client

  result = await search_tool.similarity_search(
      project_id=mock_spanner_ids["project_id"],
      instance_id=mock_spanner_ids["instance_id"],
      database_id=mock_spanner_ids["database_id"],
      table_name=mock_spanner_ids["table_name"],
      query="test query",
      embedding_column_to_search="embedding_col",
      columns=["col1"],
      embedding_options={
          "spanner_postgresql_vertex_ai_embedding_model_endpoint": (
              "test_endpoint"
          )
      },
      credentials=mock_credentials,
      search_options={
          "nearest_neighbors_algorithm": "APPROXIMATE_NEAREST_NEIGHBORS"
      },
  )
  assert result["status"] == "ERROR"
  assert (
      "APPROXIMATE_NEAREST_NEIGHBORS is not supported for PostgreSQL dialect."
      in result["error_details"]
  )


@pytest.mark.asyncio
@mock.patch.object(client, "get_spanner_client")
async def test_similarity_search_gsql_missing_embedding_model_error(
    mock_get_spanner_client, mock_spanner_ids, mock_credentials
):
  """Test similarity_search with missing embedding_options for GoogleSQL dialect."""
  mock_spanner_client = MagicMock()
  mock_instance = MagicMock()
  mock_database = MagicMock()
  mock_database.database_dialect = DatabaseDialect.GOOGLE_STANDARD_SQL
  mock_instance.database.return_value = mock_database
  mock_spanner_client.instance.return_value = mock_instance
  mock_get_spanner_client.return_value = mock_spanner_client

  result = await search_tool.similarity_search(
      project_id=mock_spanner_ids["project_id"],
      instance_id=mock_spanner_ids["instance_id"],
      database_id=mock_spanner_ids["database_id"],
      table_name=mock_spanner_ids["table_name"],
      query="test query",
      embedding_column_to_search="embedding_col",
      columns=["col1"],
      embedding_options={
          "spanner_postgresql_vertex_ai_embedding_model_endpoint": (
              "test_endpoint"
          )
      },
      credentials=mock_credentials,
  )
  assert result["status"] == "ERROR"
  assert (
      "embedding_options['vertex_ai_embedding_model_name'] or"
      " embedding_options['spanner_googlesql_embedding_model_name'] must be"
      " specified for GoogleSQL dialect Spanner database."
      in result["error_details"]
  )


@pytest.mark.asyncio
@mock.patch.object(client, "get_spanner_client")
async def test_similarity_search_pg_missing_embedding_model_error(
    mock_get_spanner_client, mock_spanner_ids, mock_credentials
):
  """Test similarity_search with missing embedding_options for PostgreSQL dialect."""
  mock_spanner_client = MagicMock()
  mock_instance = MagicMock()
  mock_database = MagicMock()
  mock_database.database_dialect = DatabaseDialect.POSTGRESQL
  mock_instance.database.return_value = mock_database
  mock_spanner_client.instance.return_value = mock_instance
  mock_get_spanner_client.return_value = mock_spanner_client

  result = await search_tool.similarity_search(
      project_id=mock_spanner_ids["project_id"],
      instance_id=mock_spanner_ids["instance_id"],
      database_id=mock_spanner_ids["database_id"],
      table_name=mock_spanner_ids["table_name"],
      query="test query",
      embedding_column_to_search="embedding_col",
      columns=["col1"],
      embedding_options={
          "spanner_googlesql_embedding_model_name": "EmbeddingsModel"
      },
      credentials=mock_credentials,
  )
  assert result["status"] == "ERROR"
  assert (
      "embedding_options['vertex_ai_embedding_model_name'] or"
      " embedding_options['spanner_postgresql_vertex_ai_embedding_model_endpoint']"
      " must be specified for PostgreSQL dialect Spanner database."
      in result["error_details"]
  )


@pytest.mark.parametrize(
    "embedding_options",
    [
        pytest.param(
            {
                "vertex_ai_embedding_model_name": "test-model",
                "spanner_googlesql_embedding_model_name": "test-model-2",
            },
            id="vertex_ai_and_googlesql",
        ),
        pytest.param(
            {
                "vertex_ai_embedding_model_name": "test-model",
                "spanner_postgresql_vertex_ai_embedding_model_endpoint": (
                    "test-endpoint"
                ),
            },
            id="vertex_ai_and_postgresql",
        ),
        pytest.param(
            {
                "spanner_googlesql_embedding_model_name": "test-model",
                "spanner_postgresql_vertex_ai_embedding_model_endpoint": (
                    "test-endpoint"
                ),
            },
            id="googlesql_and_postgresql",
        ),
        pytest.param(
            {
                "vertex_ai_embedding_model_name": "test-model",
                "spanner_googlesql_embedding_model_name": "test-model-2",
                "spanner_postgresql_vertex_ai_embedding_model_endpoint": (
                    "test-endpoint"
                ),
            },
            id="all_three_models",
        ),
        pytest.param(
            {},
            id="no_models",
        ),
    ],
)
@pytest.mark.asyncio
@mock.patch.object(client, "get_spanner_client")
async def test_similarity_search_multiple_embedding_options_error(
    mock_get_spanner_client,
    mock_spanner_ids,
    mock_credentials,
    embedding_options,
):
  """Test similarity_search with multiple embedding models."""
  mock_spanner_client = MagicMock()
  mock_instance = MagicMock()
  mock_database = MagicMock()
  mock_database.database_dialect = DatabaseDialect.GOOGLE_STANDARD_SQL
  mock_instance.database.return_value = mock_database
  mock_spanner_client.instance.return_value = mock_instance
  mock_get_spanner_client.return_value = mock_spanner_client

  result = await search_tool.similarity_search(
      project_id=mock_spanner_ids["project_id"],
      instance_id=mock_spanner_ids["instance_id"],
      database_id=mock_spanner_ids["database_id"],
      table_name=mock_spanner_ids["table_name"],
      query="test query",
      embedding_column_to_search="embedding_col",
      columns=["col1"],
      embedding_options=embedding_options,
      credentials=mock_credentials,
  )
  assert result["status"] == "ERROR"
  assert (
      "Exactly one embedding model option must be specified."
      in result["error_details"]
  )


@pytest.mark.asyncio
@mock.patch.object(client, "get_spanner_client")
async def test_similarity_search_output_dimensionality_gsql_error(
    mock_get_spanner_client, mock_spanner_ids, mock_credentials
):
  """Test similarity_search with output_dimensionality and spanner_googlesql_embedding_model_name."""
  mock_spanner_client = MagicMock()
  mock_instance = MagicMock()
  mock_database = MagicMock()
  mock_database.database_dialect = DatabaseDialect.GOOGLE_STANDARD_SQL
  mock_instance.database.return_value = mock_database
  mock_spanner_client.instance.return_value = mock_instance
  mock_get_spanner_client.return_value = mock_spanner_client

  result = await search_tool.similarity_search(
      project_id=mock_spanner_ids["project_id"],
      instance_id=mock_spanner_ids["instance_id"],
      database_id=mock_spanner_ids["database_id"],
      table_name=mock_spanner_ids["table_name"],
      query="test query",
      embedding_column_to_search="embedding_col",
      columns=["col1"],
      embedding_options={
          "spanner_googlesql_embedding_model_name": "EmbeddingsModel",
          "output_dimensionality": 128,
      },
      credentials=mock_credentials,
  )
  assert result["status"] == "ERROR"
  assert "is not supported when" in result["error_details"]


@pytest.mark.asyncio
@mock.patch.object(client, "get_spanner_client")
async def test_similarity_search_unsupported_algorithm_error(
    mock_get_spanner_client, mock_spanner_ids, mock_credentials
):
  """Test similarity_search with an unsupported nearest neighbors algorithm."""
  mock_spanner_client = MagicMock()
  mock_instance = MagicMock()
  mock_database = MagicMock()
  mock_database.database_dialect = DatabaseDialect.GOOGLE_STANDARD_SQL
  mock_instance.database.return_value = mock_database
  mock_spanner_client.instance.return_value = mock_instance
  mock_get_spanner_client.return_value = mock_spanner_client

  result = await search_tool.similarity_search(
      project_id=mock_spanner_ids["project_id"],
      instance_id=mock_spanner_ids["instance_id"],
      database_id=mock_spanner_ids["database_id"],
      table_name=mock_spanner_ids["table_name"],
      query="test query",
      embedding_column_to_search="embedding_col",
      columns=["col1"],
      embedding_options={"vertex_ai_embedding_model_name": "test-model"},
      credentials=mock_credentials,
      search_options={"nearest_neighbors_algorithm": "INVALID_ALGORITHM"},
  )
  assert result["status"] == "ERROR"
  assert "Unsupported search_options" in result["error_details"]