DocsGPT/tests/api/answer/test_stream_processor_exposure.py

"""Tests for per-source search exposure partitioning (E1 / D11)."""

from __future__ import annotations

from unittest.mock import MagicMock, patch

import pytest

from application.api.answer.services.stream_processor import StreamProcessor
from application.storage.db.source_config import RetrievalConfig


def _processor() -> StreamProcessor:
    """A bare StreamProcessor with only the fields the exposure helpers touch.

    Built via ``__new__`` to skip the DB-heavy ``__init__``; the exposure
    methods read ``all_sources`` / ``source`` / ``agent_config`` only.
    """
    sp = StreamProcessor.__new__(StreamProcessor)
    sp.all_sources = []
    sp.source = {}
    sp.agent_config = {}
    sp.retriever_config = {"retriever_name": "classic", "chunks": 2, "doc_token_limit": 50000}
    sp.data = {}
    return sp


@pytest.mark.unit
class TestExposurePartition:
    def test_no_config_all_default_to_prefetch(self):
        sp = _processor()
        sp.all_sources = [{"id": "a", "retrieval": None}, {"id": "b", "retrieval": None}]
        prefetch, agentic = sp._exposure_partition()
        assert [e["id"] for e in prefetch] == ["a", "b"]
        assert agentic == []

    def test_mixed_partition(self):
        sp = _processor()
        sp.all_sources = [
            {"id": "a", "retrieval": RetrievalConfig(exposure="prefetch")},
            {"id": "b", "retrieval": RetrievalConfig(exposure="agentic_tool")},
            {"id": "c", "retrieval": RetrievalConfig()},  # default prefetch
        ]
        prefetch, agentic = sp._exposure_partition()
        assert sorted(e["id"] for e in prefetch) == ["a", "c"]
        assert [e["id"] for e in agentic] == ["b"]

    def test_exposure_of_dict_and_model_and_missing(self):
        sp = _processor()
        assert sp._exposure_of(RetrievalConfig(exposure="agentic_tool")) == "agentic_tool"
        assert sp._exposure_of({"exposure": "agentic_tool"}) == "agentic_tool"
        assert sp._exposure_of(None) == "prefetch"
        assert sp._exposure_of({}) == "prefetch"

    def test_source_for_docs(self):
        sp = _processor()
        assert sp._source_for_docs(["a", "b"]) == {"active_docs": ["a", "b"]}
        assert sp._source_for_docs([]) == {}


@pytest.mark.unit
class TestBuildAgentExposure:
    def _agentic_processor(self):
        sp = _processor()
        sp.agent_config = {"agent_type": "agentic"}
        sp.initialize = MagicMock()
        sp.pre_fetch_docs = MagicMock(return_value=("docs_together", ["doc"]))
        sp.pre_fetch_tools = MagicMock(return_value={"tool": {}})
        sp.create_agent = MagicMock(return_value="AGENT")
        return sp

    def test_all_prefetch_is_today_no_partition(self):
        # No agentic_tool source → today's behavior: no pre-fetch, no
        # agentic_sources passed (tool exposes all sources).
        sp = self._agentic_processor()
        sp.all_sources = [
            {"id": "a", "retrieval": RetrievalConfig()},
            {"id": "b", "retrieval": None},
        ]
        result = sp.build_agent("q")
        assert result == "AGENT"
        sp.pre_fetch_docs.assert_not_called()
        _, kwargs = sp.create_agent.call_args
        assert "agentic_sources" not in kwargs

    def test_mixed_prefetches_and_scopes_tool(self):
        sp = self._agentic_processor()
        sp.all_sources = [
            {"id": "a", "retrieval": RetrievalConfig(exposure="prefetch")},
            {"id": "b", "retrieval": RetrievalConfig(exposure="agentic_tool")},
        ]
        sp.build_agent("q")
        # Pre-fetch ran scoped to the prefetch subset.
        sp.pre_fetch_docs.assert_called_once()
        _, pf_kwargs = sp.pre_fetch_docs.call_args
        assert pf_kwargs.get("exposure") == "prefetch"
        # create_agent received only the agentic_tool subset as the tool sources.
        _, kwargs = sp.create_agent.call_args
        assert [e["id"] for e in kwargs["agentic_sources"]] == ["b"]

    def _classic_processor(self):
        sp = _processor()
        sp.agent_config = {"agent_type": "classic"}
        sp.initialize = MagicMock()
        sp.pre_fetch_docs = MagicMock(return_value=("docs_together", ["doc"]))
        sp.pre_fetch_tools = MagicMock(return_value=None)
        sp.create_agent = MagicMock(return_value="AGENT")
        return sp

    def test_classic_default_prefetches_all_no_partition(self):
        # Default classic (all prefetch / no config): byte-identical to today —
        # unscoped pre-fetch and no agentic_sources, so no search tool is added.
        sp = self._classic_processor()
        sp.all_sources = [
            {"id": "a", "retrieval": RetrievalConfig()},
            {"id": "b", "retrieval": None},
        ]
        result = sp.build_agent("q")
        assert result == "AGENT"
        sp.pre_fetch_docs.assert_called_once_with("q")
        _, kwargs = sp.create_agent.call_args
        assert "agentic_sources" not in kwargs

    def test_classic_no_per_source_detail_is_today(self):
        # Single-source / no-config requests carry no per-source detail; classic
        # must fall back to the unscoped pre-fetch (no exposure scoping).
        sp = self._classic_processor()
        sp.all_sources = []
        sp.source = {"active_docs": ["a"]}
        sp.build_agent("q")
        sp.pre_fetch_docs.assert_called_once_with("q")
        _, kwargs = sp.create_agent.call_args
        assert "agentic_sources" not in kwargs

    def test_classic_mixed_prefetches_and_scopes_tool(self):
        # Classic with an agentic_tool source now pre-fetches only the prefetch
        # subset and exposes the agentic_tool subset via the search tool.
        sp = self._classic_processor()
        sp.all_sources = [
            {"id": "a", "retrieval": RetrievalConfig(exposure="prefetch")},
            {"id": "b", "retrieval": RetrievalConfig(exposure="agentic_tool")},
        ]
        sp.build_agent("q")
        sp.pre_fetch_docs.assert_called_once()
        _, pf_kwargs = sp.pre_fetch_docs.call_args
        assert pf_kwargs.get("exposure") == "prefetch"
        _, kwargs = sp.create_agent.call_args
        assert [e["id"] for e in kwargs["agentic_sources"]] == ["b"]


@pytest.mark.unit
class TestNonAgentSourceConfig:
    """The non-agent chat path must also load per-source config so exposure
    (and other per-source overrides) are honored, not just the agent path."""

    def test_active_docs_loads_per_source_retrieval_config(self):
        sp = StreamProcessor.__new__(StreamProcessor)
        sp._agent_data = None
        sp.data = {"active_docs": "s1"}
        sp.initial_user_id = "user1"
        sp.source = {}
        sp.all_sources = []
        fake_source = {
            "config": {"retrieval": {"exposure": "agentic_tool", "chunks": 7}}
        }
        with patch(
            "application.api.answer.services.stream_processor.db_readonly"
        ), patch(
            "application.api.answer.services.stream_processor.can_access",
            return_value=True,
        ), patch(
            "application.api.answer.services.stream_processor.SourcesRepository"
        ) as repo:
            # Read unscoped after the access check, so a team grantee gets the
            # source's real config instead of silently falling back to defaults.
            repo.return_value.get_by_id.return_value = fake_source
            sp._configure_source()
        assert sp.source == {"active_docs": "s1"}
        assert len(sp.all_sources) == 1
        assert sp.all_sources[0]["id"] == "s1"
        assert sp.all_sources[0]["retrieval"].exposure == "agentic_tool"
        assert sp.all_sources[0]["retrieval"].chunks == 7