DocsGPT/tests/agents/tools/test_internal_search_pg.py

"""Real-PG tests for internal_search directory-structure loading."""

from contextlib import contextmanager
from unittest.mock import patch



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

    with patch(
        "application.storage.db.session.db_readonly", _yield
    ):
        yield


class TestGetDirectoryStructureFromDb:
    def test_returns_none_for_empty_active_docs(self):
        from application.agents.tools.internal_search import InternalSearchTool

        tool = InternalSearchTool({"source": {}})
        assert tool._get_directory_structure() is None

    def test_loads_single_source_structure(self, pg_conn):
        from application.agents.tools.internal_search import InternalSearchTool
        from application.storage.db.repositories.sources import (
            SourcesRepository,
        )

        user = "u-ds"
        src = SourcesRepository(pg_conn).create(
            "hello",
            user_id=user,
            directory_structure={"file.txt": {"type": "file", "size_bytes": 10}},
        )
        tool = InternalSearchTool({
            "source": {"active_docs": str(src["id"])},
            "decoded_token": {"sub": user},
        })

        with _patch_db(pg_conn):
            got = tool._get_directory_structure()
        assert got is not None
        assert "file.txt" in got

    def test_loads_multiple_sources_merges(self, pg_conn):
        from application.agents.tools.internal_search import InternalSearchTool
        from application.storage.db.repositories.sources import (
            SourcesRepository,
        )

        user = "u-ds-multi"
        repo = SourcesRepository(pg_conn)
        src1 = repo.create(
            "src1", user_id=user,
            directory_structure={"a.txt": {"type": "file", "size_bytes": 5}},
        )
        src2 = repo.create(
            "src2", user_id=user,
            directory_structure={"b.txt": {"type": "file", "size_bytes": 5}},
        )
        tool = InternalSearchTool({
            "source": {"active_docs": [str(src1["id"]), str(src2["id"])]},
            "decoded_token": {"sub": user},
        })

        with _patch_db(pg_conn):
            got = tool._get_directory_structure()
        assert got is not None
        # Multi-source merging uses source name as key
        assert "src1" in got or "a.txt" in got

    def test_skips_missing_source(self, pg_conn):
        from application.agents.tools.internal_search import InternalSearchTool

        tool = InternalSearchTool({
            "source": {"active_docs": "00000000-0000-0000-0000-000000000000"},
            "decoded_token": {"sub": "u"},
        })
        with _patch_db(pg_conn):
            got = tool._get_directory_structure()
        assert got is None


class TestSourcesHaveDirectoryStructureDb:
    def test_false_no_active_docs(self):
        from application.agents.tools.internal_search import (
            sources_have_directory_structure,
        )
        assert sources_have_directory_structure({}) is False

    def test_true_when_source_has_structure(self, pg_conn):
        from application.agents.tools.internal_search import (
            sources_have_directory_structure,
        )
        from application.storage.db.repositories.sources import (
            SourcesRepository,
        )

        src = SourcesRepository(pg_conn).create(
            "s", user_id="u",
            directory_structure={"a.txt": {"type": "file", "size_bytes": 5}},
        )
        with _patch_db(pg_conn):
            got = sources_have_directory_structure(
                {"active_docs": str(src["id"])}
            )
        assert got is True

    def test_false_when_source_has_no_structure(self, pg_conn):
        from application.agents.tools.internal_search import (
            sources_have_directory_structure,
        )
        from application.storage.db.repositories.sources import (
            SourcesRepository,
        )

        src = SourcesRepository(pg_conn).create("s", user_id="u")
        with _patch_db(pg_conn):
            got = sources_have_directory_structure(
                {"active_docs": str(src["id"])}
            )
        assert got is False

    def test_legacy_id_lookup(self, pg_conn):
        from application.agents.tools.internal_search import (
            sources_have_directory_structure,
        )
        from application.storage.db.repositories.sources import (
            SourcesRepository,
        )

        SourcesRepository(pg_conn).create(
            "s", user_id="u",
            legacy_mongo_id="507f1f77bcf86cd799439011",
            directory_structure={"x.txt": {"type": "file", "size_bytes": 5}},
        )
        with _patch_db(pg_conn):
            got = sources_have_directory_structure(
                {"active_docs": "507f1f77bcf86cd799439011"}
            )
        assert got is True