DocsGPT/tests/connectors/test_service_read.py

"""Tests for the read side of the connection service and its API."""

from __future__ import annotations

import json
from contextlib import contextmanager
from unittest.mock import patch

import pytest
from flask import Flask
from sqlalchemy import text

from docsgpt.connectors import service


def _session(conn, user="alice", provider="google_drive", **cols) -> str:
    cols = {"status": "authorized", "user_email": f"{user}@example.com", **cols}
    casts = {"token_info": "jsonb", "session_data": "jsonb"}
    names = ", ".join(["user_id", "provider", *cols])
    values = ", ".join(
        [":user_id", ":provider", *[f"CAST(:{k} AS {casts[k]})" if k in casts else f":{k}" for k in cols]]
    )
    return str(
        conn.execute(
            text(f"INSERT INTO connector_sessions ({names}) VALUES ({values}) RETURNING id"),
            {"user_id": user, "provider": provider, **cols},
        ).scalar()
    )


def _source(conn, connection_id, user="alice", name="Handbook") -> str:
    return str(
        conn.execute(
            text(
                "INSERT INTO sources (user_id, name, type, sync_frequency, connection_id) "
                "VALUES (:u, :n, 'connector:file', 'weekly', CAST(:c AS uuid)) RETURNING id"
            ),
            {"u": user, "n": name, "c": connection_id},
        ).scalar()
    )


def _tool(conn, connection_id, user="alice") -> str:
    actions = [
        {"name": "search_pages", "description": "Search", "active": True},
        {"name": "create_page", "description": "Create", "active": True, "require_approval": True},
    ]
    return str(
        conn.execute(
            text(
                "INSERT INTO user_tools (user_id, name, display_name, actions, connection_id) "
                "VALUES (:u, 'mcp_tool', 'Notion', CAST(:a AS jsonb), CAST(:c AS uuid)) RETURNING id"
            ),
            {"u": user, "a": json.dumps(actions), "c": connection_id},
        ).scalar()
    )


class TestNormalizeStatus:
    @pytest.mark.parametrize(
        "row,expected",
        [
            ({"status": "authorized"}, "connected"),
            ({"status": "reconnect_needed"}, "reconnect_needed"),
            ({"status": "pending"}, "pending"),
            ({"status": "pending", "token_info": {"access_token": "x"}}, "connected"),
            ({"status": None, "session_data": {"tokens": {"access_token": "x"}}}, "connected"),
            ({"status": None, "session_data": {"client_info": {}}}, "pending"),
            ({"status": None, "encrypted_credentials": "v2:..."}, "connected"),
        ],
    )
    def test_statuses(self, row, expected):
        assert service.normalize_status(row) == expected

    def test_worst_status(self):
        assert service.worst_status(["connected", "reconnect_needed"]) == "reconnect_needed"
        assert service.worst_status([]) is None


class TestListing:
    def test_lists_finished_connections_with_counts(self, pg_conn):
        drive = _session(pg_conn)
        _session(pg_conn, provider="confluence", status="pending", user_email=None)
        _source(pg_conn, drive)
        _source(pg_conn, drive, name="Wiki")
        connections = service.list_connections(pg_conn, "alice")
        assert [c["connector_key"] for c in connections] == ["google_drive"]
        only = connections[0]
        assert only["name"] == "Google Drive"
        assert only["account_label"] == "alice@example.com"
        assert only["status"] == "connected"
        assert only["source_count"] == 2
        assert "token_info" not in only and "session_token" not in only

    def test_other_users_rows_hidden(self, pg_conn):
        _session(pg_conn, user="bob")
        assert service.list_connections(pg_conn, "alice") == []

    def test_detail_lists_sources_and_tools(self, pg_conn):
        from docsgpt.storage.db.repositories.connector_sessions import ConnectorSessionsRepository

        drive = _session(pg_conn, status="reconnect_needed")
        _source(pg_conn, drive)
        _tool(pg_conn, drive)
        row = ConnectorSessionsRepository(pg_conn).get(drive)
        detail = service.connection_detail(pg_conn, row)
        assert detail["sources"][0]["sync_state"] == "paused_reconnect"
        actions = {a["name"]: a for a in detail["tools"][0]["actions"]}
        assert actions["search_pages"] == {
            "name": "search_pages", "description": "Search", "access": "read", "permission": "always",
            "parameters": [],
        }
        assert actions["create_page"]["access"] == "write"
        assert actions["create_page"]["permission"] == "ask"


class TestCatalogForUser:
    def test_states(self, pg_conn, monkeypatch):
        from docsgpt.core.settings import settings

        monkeypatch.setattr(settings, "GOOGLE_CLIENT_ID", "id")
        monkeypatch.setattr(settings, "GOOGLE_CLIENT_SECRET", "secret")
        monkeypatch.setattr(settings, "MICROSOFT_CLIENT_ID", None)
        _session(pg_conn)
        _session(pg_conn, provider="confluence", status="reconnect_needed")
        entries = {e["key"]: e for e in service.catalog_for_user(pg_conn, "alice", is_admin=False)}
        assert entries["google_drive"]["state"] == "connected"
        assert entries["google_drive"]["connected_count"] == 1
        assert entries["confluence"]["state"] == "reconnect"
        # Needs server settings and nobody connected it: members never see it.
        assert "share_point" not in entries
        assert entries["telegram"]["state"] == "available"
        assert entries["custom_mcp"]["state"] == "custom"

    def test_unconfigured_connector_stays_hidden_even_when_switched_on(self, pg_conn, monkeypatch):
        from docsgpt.core.settings import settings

        monkeypatch.setattr(settings, "MICROSOFT_CLIENT_ID", None)
        entries = {
            e["key"]: e
            for e in service.catalog_for_user(
                pg_conn, "alice", is_admin=True, policies={"share_point": {"enabled": True}}
            )
        }
        assert "share_point" not in entries

    def test_policy_hides_a_disabled_connector(self, pg_conn):
        entries = {
            e["key"]: e
            for e in service.catalog_for_user(
                pg_conn, "alice", is_admin=False, policies={"telegram": {"enabled": False}}
            )
        }
        assert "telegram" not in entries

    def test_disabled_connector_with_a_connection_shows_as_turned_off(self, pg_conn):
        """Members can still see and remove what they connected before it was turned off."""
        _session(pg_conn, provider="telegram")
        entries = {
            e["key"]: e
            for e in service.catalog_for_user(
                pg_conn, "alice", is_admin=False, policies={"telegram": {"enabled": False}}
            )
        }
        assert entries["telegram"]["state"] == "disabled"
        assert entries["telegram"]["available"] is False


class TestConnectorIsEnabled:
    def test_unconfigured_connector_is_off_by_default(self, monkeypatch):
        from docsgpt.core.settings import settings

        monkeypatch.setattr(settings, "GOOGLE_CLIENT_ID", None)
        assert service.connector_is_enabled({}, "google_drive") is False
        assert service.connector_is_enabled({}, "telegram") is True

    def test_configured_connector_is_on_by_default(self, monkeypatch):
        from docsgpt.core.settings import settings

        monkeypatch.setattr(settings, "GOOGLE_CLIENT_ID", "id")
        monkeypatch.setattr(settings, "GOOGLE_CLIENT_SECRET", "secret")
        assert service.connector_is_enabled({}, "google_drive") is True

    def test_an_explicit_switch_wins(self, monkeypatch):
        from docsgpt.core.settings import settings

        monkeypatch.setattr(settings, "GOOGLE_CLIENT_ID", None)
        assert service.connector_is_enabled({"google_drive": {"enabled": True}}, "google_drive") is True
        assert service.connector_is_enabled({"telegram": {"enabled": False}}, "telegram") is False
        assert service.connector_is_enabled({"telegram": {"enabled": None}}, "telegram") is True


@contextmanager
def _patched_db(conn):
    @contextmanager
    def _yield():
        yield conn

    with patch("docsgpt.api.connector.connections.db_readonly", _yield):
        yield


@pytest.fixture
def app():
    return Flask(__name__)


def _call(app, resource, path, token, *args):
    with app.test_request_context(path):
        from flask import request

        request.decoded_token = token
        return resource().get(*args)


class TestRoutes:
    def test_catalog_requires_auth(self, app):
        from docsgpt.api.connector.connections import ConnectorCatalog

        assert _call(app, ConnectorCatalog, "/api/connectors/catalog", None).status_code == 401

    def test_catalog(self, app, pg_conn):
        from docsgpt.api.connector.connections import ConnectorCatalog

        with _patched_db(pg_conn):
            resp = _call(app, ConnectorCatalog, "/api/connectors/catalog", {"sub": "alice"})
        assert resp.status_code == 200
        keys = {c["key"] for c in resp.get_json()["connectors"]}
        assert "telegram" in keys

    def test_list_and_detail(self, app, pg_conn):
        from docsgpt.api.connector.connections import ConnectionDetail, ConnectionsList

        drive = _session(pg_conn)
        with _patched_db(pg_conn):
            listed = _call(app, ConnectionsList, "/api/connections", {"sub": "alice"})
            detail = _call(app, ConnectionDetail, f"/api/connections/{drive}", {"sub": "alice"}, drive)
        assert listed.get_json()["connections"][0]["id"] == drive
        assert detail.get_json()["connection"]["id"] == drive

    def test_detail_of_another_users_connection_is_404(self, app, pg_conn):
        from docsgpt.api.connector.connections import ConnectionDetail

        drive = _session(pg_conn, user="bob")
        with _patched_db(pg_conn):
            resp = _call(app, ConnectionDetail, f"/api/connections/{drive}", {"sub": "alice"}, drive)
        assert resp.status_code == 404

    def test_detail_of_bad_id_is_404(self, app, pg_conn):
        from docsgpt.api.connector.connections import ConnectionDetail

        with _patched_db(pg_conn):
            resp = _call(app, ConnectionDetail, "/api/connections/nope", {"sub": "alice"}, "nope")
        assert resp.status_code == 404


class TestDisconnect:
    def test_clears_credentials_keeps_resources(self, pg_conn):
        from docsgpt.storage.db.repositories.connector_sessions import ConnectorSessionsRepository

        drive = _session(
            pg_conn, session_token="tok", token_info=json.dumps({"access_token": "at", "refresh_token": "rt"}),
        )
        source = _source(pg_conn, drive)
        repo = ConnectorSessionsRepository(pg_conn)
        result = service.disconnect(pg_conn, repo.get(drive))
        assert result["status"] == "disconnected"
        row = repo.get(drive)
        assert row["token_info"] is None and row["session_token"] is None
        linked = pg_conn.execute(
            text("SELECT connection_id FROM sources WHERE id = CAST(:id AS uuid)"), {"id": source}
        ).scalar()
        assert str(linked) == drive

    def test_mcp_keeps_client_registration(self, pg_conn):
        from docsgpt.storage.db.repositories.connector_sessions import ConnectorSessionsRepository

        mcp = _session(
            pg_conn,
            provider="mcp:https://mcp.example.com",
            status=None,
            session_data=json.dumps({"tokens": {"access_token": "x"}, "client_info": {"client_id": "c"}}),
        )
        repo = ConnectorSessionsRepository(pg_conn)
        service.disconnect(pg_conn, repo.get(mcp))
        row = repo.get(mcp)
        assert row["session_data"] == {}
        assert service.read_secrets(row) == {"client_info": {"client_id": "c"}}

    def test_route_rejects_other_users(self, app, pg_conn):
        from docsgpt.api.connector.connections import ConnectionDisconnect

        drive = _session(pg_conn, user="bob")

        @contextmanager
        def _yield():
            yield pg_conn

        with patch("docsgpt.api.connector.connections.db_session", _yield), app.test_request_context(
            f"/api/connections/{drive}/disconnect", method="POST"
        ):
            from flask import request

            request.decoded_token = {"sub": "alice"}
            resp = ConnectionDisconnect().post(drive)
        assert resp.status_code == 404