DocsGPT/tests/worker/test_reingest_source_worker.py

"""Smoke test for ``application.worker.reingest_source_worker``.

The task reads a source row, diffs its stored ``directory_structure``
against what's currently in storage, updates the vector store, then
writes the refreshed ``directory_structure`` / ``date`` / ``tokens``
back to the ``sources`` row. We assert that last PG update actually
lands on the row for our ephemeral DB.
"""

from __future__ import annotations

from io import BytesIO
from unittest.mock import MagicMock

import pytest

from application.parser.schema.base import Document
from application.storage.db.repositories.sources import SourcesRepository


@pytest.mark.unit
class TestReingestSourceWorker:
    def test_updates_source_directory_structure_and_tokens(
        self, pg_conn, patch_worker_db, task_self, monkeypatch
    ):
        from application import worker

        # Seed a source we can re-ingest.
        src = SourcesRepository(pg_conn).create(
            "doc-set",
            user_id="alice",
            type="local",
            retriever="classic",
            file_path="inputs/alice/doc-set",
            directory_structure={
                "stale.txt": {
                    "type": "text/plain",
                    "size_bytes": 10,
                    "token_count": 5,
                }
            },
        )
        source_id = str(src["id"])

        # Storage: pretend the source directory has one file now, a
        # different one from what the source row last stored — that way
        # the "added/removed" branch runs end-to-end.
        fake_storage = MagicMock(name="storage")
        fake_storage.is_directory.return_value = True
        fake_storage.list_files.return_value = []
        monkeypatch.setattr(
            worker.StorageCreator, "get_storage", lambda: fake_storage
        )

        # Reader: return no docs but advertise the new directory shape and
        # per-file token counts. The worker's update statement uses
        # ``sum(reader.file_token_counts.values())`` and the ``
        # directory_structure`` as-is.
        fake_reader = MagicMock(name="reader")
        fake_reader.load_data.return_value = []
        fake_reader.directory_structure = {
            "fresh.md": {
                "type": "text/markdown",
                "size_bytes": 42,
                "token_count": 17,
            }
        }
        fake_reader.file_token_counts = {"fresh.md": 17}
        monkeypatch.setattr(
            worker, "SimpleDirectoryReader", lambda *a, **kw: fake_reader
        )

        # Vector store: report no existing chunks, and swallow any add/
        # delete call.
        fake_store = MagicMock(name="vector_store")
        fake_store.get_chunks.return_value = []
        monkeypatch.setattr(
            "application.vectorstore.vector_creator.VectorCreator.create_vectorstore",
            lambda *a, **kw: fake_store,
        )

        result = worker.reingest_source_worker(task_self, source_id, "alice")

        assert result["status"] == "completed"
        # ``fresh.md`` is new vs ``stale.txt``, which was removed.
        assert "fresh.md" in result["added_files"]
        assert "stale.txt" in result["removed_files"]

        refreshed = SourcesRepository(pg_conn).get(source_id, "alice")
        assert refreshed is not None
        # Token count was recomputed from the reader.
        assert refreshed["tokens"] == "17"
        # And the new directory_structure replaced the stale one.
        assert "fresh.md" in refreshed["directory_structure"]
        assert "stale.txt" not in refreshed["directory_structure"]

    def test_reingest_threads_source_chunking_config(
        self, pg_conn, patch_worker_db, task_self, monkeypatch
    ):
        """The source's stored ``config.chunking`` must drive re-chunking.

        Seeds a source with a non-default chunking config and an added file,
        then asserts ``ChunkerCreator.create_chunker`` is built with that
        config — not the 1250/150 classic defaults (D1/D8).
        """
        from application import worker

        src = SourcesRepository(pg_conn).create(
            "doc-set",
            user_id="alice",
            type="local",
            retriever="classic",
            file_path="inputs/alice/doc-set",
            config={
                "chunking": {
                    "strategy": "recursive",
                    "max_tokens": 800,
                    "min_tokens": 50,
                }
            },
            directory_structure={},
        )
        source_id = str(src["id"])

        # Storage: one file in the source directory, downloaded to temp_dir so
        # the worker's added-files branch finds it on disk and chunks it.
        fake_storage = MagicMock(name="storage")
        fake_storage.is_directory.side_effect = lambda p: p == "inputs/alice/doc-set"
        fake_storage.list_files.return_value = ["inputs/alice/doc-set/fresh.md"]
        fake_storage.get_file.return_value = BytesIO(b"fresh body")
        monkeypatch.setattr(
            worker.StorageCreator, "get_storage", lambda: fake_storage
        )

        # Reader: advertise the new file and hand back one doc to chunk.
        fake_reader = MagicMock(name="reader")
        fake_reader.load_data.return_value = [
            Document(text="fresh body", extra_info={"source": "fresh.md"})
        ]
        fake_reader.directory_structure = {
            "fresh.md": {
                "type": "text/markdown",
                "size_bytes": 10,
                "token_count": 3,
            }
        }
        fake_reader.file_token_counts = {"fresh.md": 3}
        monkeypatch.setattr(
            worker, "SimpleDirectoryReader", lambda *a, **kw: fake_reader
        )

        fake_store = MagicMock(name="vector_store")
        fake_store.get_chunks.return_value = []
        monkeypatch.setattr(
            "application.vectorstore.vector_creator.VectorCreator.create_vectorstore",
            lambda *a, **kw: fake_store,
        )

        calls: list[dict] = []

        def _create_chunker(strategy, **kwargs):
            calls.append({"strategy": strategy, **kwargs})
            chunker = MagicMock(name="chunker")
            chunker.chunk.side_effect = lambda documents: documents
            return chunker

        monkeypatch.setattr(
            worker.ChunkerCreator,
            "create_chunker",
            staticmethod(_create_chunker),
        )

        result = worker.reingest_source_worker(task_self, source_id, "alice")

        assert result["status"] == "completed"
        assert "fresh.md" in result["added_files"]
        assert len(calls) == 1
        # Config-driven, not the 1250/150 classic defaults.
        assert calls[0]["strategy"] == "recursive"
        assert calls[0]["chunking_strategy"] == "recursive"
        assert calls[0]["max_tokens"] == 800
        assert calls[0]["min_tokens"] == 50