test(integration): enhance Drive indexer credential resolution tests for Composio and native connectors

surfsense_backend/tests/integration/google_unification/test_drive_indexer_credentials.py

@@ -1,13 +1,26 @@
-"""Integration tests: Drive indexer credential resolution for Composio vs native connectors.
+"""Integration tests: Drive indexer client + credential resolution.
 
-Exercises ``index_google_drive_files`` with a real PostgreSQL database
-containing seeded connector records.  Google API and Composio SDK are
-mocked at their system boundaries.
+Locks in the post-cea8618 architectural contract:
+
+- Composio Drive connectors MUST use ``ComposioDriveClient`` (which routes
+  through ``composio.tools.execute``) and MUST NOT depend on a raw OAuth
+  access token via ``ComposioService.get_access_token``.
+- Native Drive connectors MUST continue to use ``GoogleDriveClient`` with
+  credentials loaded from the connector config (no Composio involvement).
+- Composio Drive connectors missing ``composio_connected_account_id`` MUST
+  short-circuit with a clear error before any client is constructed.
+
+Background: prior to ``cea8618`` the Composio path used
+``build_composio_credentials → GoogleDriveClient``. That broke in production
+once Composio's "Mask Connected Account Secrets" project toggle was on,
+because the masked token failed the ``len(access_token) >= 20`` guard in
+``ComposioService.get_access_token``. The structural assertions here make
+any future regression to that token-based path fail at PR time.
 """
 
 from __future__ import annotations
 
-from unittest.mock import MagicMock, patch
+from unittest.mock import ANY, AsyncMock, MagicMock, patch
 
 import pytest
 import pytest_asyncio
@@ -25,6 +38,21 @@ pytestmark = pytest.mark.integration
 
 _COMPOSIO_ACCOUNT_ID = "composio-test-account-123"
 _INDEXER_MODULE = "app.tasks.connector_indexers.google_drive_indexer"
+_GET_ACCESS_TOKEN = "app.services.composio_service.ComposioService.get_access_token"
+
+
+def _mock_drive_client(
+    *, list_files_return: tuple = ([], None, None)
+) -> MagicMock:
+    """Duck-typed client mock whose ``list_files`` yields the supplied tuple.
+
+    Returning an empty file list short-circuits the indexer's full-scan
+    loop after the first page so the test exercises only the
+    construction + listing path, not download / ETL / DB writes.
+    """
+    mock = MagicMock()
+    mock.list_files = AsyncMock(return_value=list_files_return)
+    return mock
 
 
 @pytest_asyncio.fixture
@@ -69,25 +97,31 @@ async def committed_composio_no_account_id(async_engine):
     await cleanup_space(async_engine, data["search_space_id"])
 
 
+@patch(_GET_ACCESS_TOKEN)
 @patch(f"{_INDEXER_MODULE}.TaskLoggingService")
 @patch(f"{_INDEXER_MODULE}.GoogleDriveClient")
-@patch(f"{_INDEXER_MODULE}.build_composio_credentials")
-async def test_composio_connector_uses_composio_credentials(
-    mock_build_creds,
-    mock_client_cls,
+@patch(f"{_INDEXER_MODULE}.ComposioDriveClient")
+async def test_composio_drive_indexer_uses_composio_drive_client(
+    mock_composio_client_cls,
+    mock_native_client_cls,
     mock_task_logger_cls,
+    mock_get_access_token,
     async_engine,
     committed_drive_connector,
 ):
-    """Drive indexer calls build_composio_credentials for a Composio connector
-    and passes the result to GoogleDriveClient."""
+    """Composio Drive must construct ComposioDriveClient and never read raw tokens.
+
+    Reverting to the pre-cea8618 ``build_composio_credentials → GoogleDriveClient``
+    path would either trip ``mock_native_client_cls.assert_not_called()`` (because
+    GoogleDriveClient would be constructed) or ``mock_get_access_token.assert_not_called()``
+    (because the credential builder reads the raw token).
+    """
     from app.tasks.connector_indexers.google_drive_indexer import (
         index_google_drive_files,
     )
 
     data = committed_drive_connector
-    mock_creds = MagicMock(name="composio-credentials")
-    mock_build_creds.return_value = mock_creds
+    mock_composio_client_cls.return_value = _mock_drive_client()
     mock_task_logger_cls.return_value = mock_task_logger()
 
     maker = make_session_factory(async_engine)
@@ -100,21 +134,29 @@ async def test_composio_connector_uses_composio_credentials(
             folder_id="test-folder-id",
         )
 
-    mock_build_creds.assert_called_once_with(_COMPOSIO_ACCOUNT_ID)
-    mock_client_cls.assert_called_once()
-    _, kwargs = mock_client_cls.call_args
-    assert kwargs.get("credentials") is mock_creds
+    mock_composio_client_cls.assert_called_once_with(
+        ANY,
+        data["connector_id"],
+        _COMPOSIO_ACCOUNT_ID,
+        entity_id=f"surfsense_{data['user_id']}",
+    )
+    mock_native_client_cls.assert_not_called()
+    mock_get_access_token.assert_not_called()
 
 
+@patch(_GET_ACCESS_TOKEN)
 @patch(f"{_INDEXER_MODULE}.TaskLoggingService")
-@patch(f"{_INDEXER_MODULE}.build_composio_credentials")
+@patch(f"{_INDEXER_MODULE}.GoogleDriveClient")
+@patch(f"{_INDEXER_MODULE}.ComposioDriveClient")
 async def test_composio_connector_without_account_id_returns_error(
-    mock_build_creds,
+    mock_composio_client_cls,
+    mock_native_client_cls,
     mock_task_logger_cls,
+    mock_get_access_token,
     async_engine,
     committed_composio_no_account_id,
 ):
-    """Drive indexer returns an error when Composio connector lacks connected_account_id."""
+    """Missing ``composio_connected_account_id`` must short-circuit before any client construction."""
     from app.tasks.connector_indexers.google_drive_indexer import (
         index_google_drive_files,
     )
@@ -134,28 +176,32 @@ async def test_composio_connector_without_account_id_returns_error(
 
     assert count == 0
     assert error is not None
-    assert (
-        "composio_connected_account_id" in error.lower() or "composio" in error.lower()
-    )
-    mock_build_creds.assert_not_called()
+    assert "composio" in error.lower()
+    assert "connected_account_id" in error.lower()
+    mock_composio_client_cls.assert_not_called()
+    mock_native_client_cls.assert_not_called()
+    mock_get_access_token.assert_not_called()
 
 
+@patch(_GET_ACCESS_TOKEN)
 @patch(f"{_INDEXER_MODULE}.TaskLoggingService")
+@patch(f"{_INDEXER_MODULE}.ComposioDriveClient")
 @patch(f"{_INDEXER_MODULE}.GoogleDriveClient")
-@patch(f"{_INDEXER_MODULE}.build_composio_credentials")
-async def test_native_connector_does_not_use_composio_credentials(
-    mock_build_creds,
-    mock_client_cls,
+async def test_native_connector_uses_google_drive_client(
+    mock_native_client_cls,
+    mock_composio_client_cls,
     mock_task_logger_cls,
+    mock_get_access_token,
     async_engine,
     committed_native_drive_connector,
 ):
-    """Drive indexer does NOT call build_composio_credentials for a native connector."""
+    """Native Drive connector must use GoogleDriveClient (no Composio involvement at all)."""
     from app.tasks.connector_indexers.google_drive_indexer import (
         index_google_drive_files,
     )
 
     data = committed_native_drive_connector
+    mock_native_client_cls.return_value = _mock_drive_client()
     mock_task_logger_cls.return_value = mock_task_logger()
 
     maker = make_session_factory(async_engine)
@@ -168,7 +214,6 @@ async def test_native_connector_does_not_use_composio_credentials(
             folder_id="test-folder-id",
         )
 
-    mock_build_creds.assert_not_called()
-    mock_client_cls.assert_called_once()
-    _, kwargs = mock_client_cls.call_args
-    assert kwargs.get("credentials") is None
+    mock_native_client_cls.assert_called_once_with(ANY, data["connector_id"])
+    mock_composio_client_cls.assert_not_called()
+    mock_get_access_token.assert_not_called()