DocsGPT/tests/worker/test_extract_graph.py

"""Tests for ``docsgpt.worker.extract_graph_worker``.

The worker loads the source row, fetches its chunks from the vector store, and
delegates to ``extract_graph_for_source``. ``graphrag_available``,
``VectorCreator.create_vectorstore``, and the extraction pipeline are mocked so
no live store access or LLM/model calls run; the ``sources`` row is real so
``SourcesRepository.get_any`` resolves.
"""

from __future__ import annotations

from unittest.mock import MagicMock

import pytest

from docsgpt.storage.db.repositories.sources import SourcesRepository


def _seed_source(pg_conn, config=None):
    src = SourcesRepository(pg_conn).create(
        "graph-set",
        user_id="alice",
        type="file",
        retriever="classic",
        config=config or {"kind": "graphrag", "retrieval": {"retriever": "graphrag"}},
    )
    return str(src["id"])


def _patch_store(monkeypatch, chunks):
    store = MagicMock(name="vectorstore")
    store.get_chunks.return_value = chunks
    monkeypatch.setattr(
        "docsgpt.vectorstore.vector_creator.VectorCreator.create_vectorstore",
        lambda *a, **kw: store,
    )
    return store


@pytest.mark.unit
class TestExtractGraphWorker:
    def test_fetches_chunks_and_calls_extraction(
        self, pg_conn, patch_worker_db, task_self, monkeypatch
    ):
        from docsgpt import worker

        source_id = _seed_source(pg_conn)
        chunks = [
            {"doc_id": "c1", "text": "alpha"},
            {"doc_id": "c2", "text": "beta"},
        ]
        store = _patch_store(monkeypatch, chunks)

        monkeypatch.setattr(
            "docsgpt.graphrag.graphrag_available", lambda: True
        )
        extract = MagicMock(
            name="extract_graph_for_source",
            return_value={"nodes": 3, "edges": 2, "chunks_processed": 2},
        )
        monkeypatch.setattr(
            "docsgpt.graphrag.extraction.extract_graph_for_source", extract
        )

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

        store.get_chunks.assert_called_once()
        extract.assert_called_once()
        assert extract.call_args.args[0] == source_id
        assert extract.call_args.args[1] == "alice"
        assert extract.call_args.args[2] == chunks
        assert extract.call_args.kwargs["config"].kind == "graphrag"
        assert result == {"nodes": 3, "edges": 2, "chunks_processed": 2}

    def test_unavailable_returns_status_no_extraction(
        self, pg_conn, patch_worker_db, task_self, monkeypatch
    ):
        from docsgpt import worker

        store = _patch_store(monkeypatch, [])
        monkeypatch.setattr(
            "docsgpt.graphrag.graphrag_available", lambda: False
        )
        extract = MagicMock(name="extract_graph_for_source")
        monkeypatch.setattr(
            "docsgpt.graphrag.extraction.extract_graph_for_source", extract
        )

        result = worker.extract_graph_worker(task_self, "src-x", "alice")

        assert result == {"status": "unavailable"}
        store.get_chunks.assert_not_called()
        extract.assert_not_called()

    def test_empty_chunks_still_calls_extraction(
        self, pg_conn, patch_worker_db, task_self, monkeypatch
    ):
        from docsgpt import worker

        source_id = _seed_source(pg_conn)
        _patch_store(monkeypatch, [])
        monkeypatch.setattr(
            "docsgpt.graphrag.graphrag_available", lambda: True
        )
        extract = MagicMock(
            name="extract_graph_for_source",
            return_value={"nodes": 0, "edges": 0, "chunks_processed": 0},
        )
        monkeypatch.setattr(
            "docsgpt.graphrag.extraction.extract_graph_for_source", extract
        )

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

        extract.assert_called_once()
        assert extract.call_args.args[2] == []
        assert result["chunks_processed"] == 0

    def test_publishes_completed_event(
        self, pg_conn, patch_worker_db, task_self, monkeypatch
    ):
        from docsgpt import worker

        source_id = _seed_source(pg_conn)
        _patch_store(monkeypatch, [{"doc_id": "c1", "text": "alpha"}])
        monkeypatch.setattr(
            "docsgpt.graphrag.graphrag_available", lambda: True
        )
        monkeypatch.setattr(
            "docsgpt.graphrag.extraction.extract_graph_for_source",
            MagicMock(return_value={"nodes": 1, "edges": 0, "chunks_processed": 1}),
        )
        events = []
        monkeypatch.setattr(
            worker, "publish_user_event",
            lambda user, etype, payload, **kw: events.append((etype, payload)),
        )

        worker.extract_graph_worker(task_self, source_id, "alice")

        types = [e[0] for e in events]
        assert "graph.extract.progress" in types
        assert types[-1] == "graph.extract.completed"
        assert events[-1][1]["nodes"] == 1

    def test_publishes_failed_event_on_error(
        self, pg_conn, patch_worker_db, task_self, monkeypatch
    ):
        from docsgpt import worker

        source_id = _seed_source(pg_conn)
        _patch_store(monkeypatch, [{"doc_id": "c1", "text": "alpha"}])
        monkeypatch.setattr(
            "docsgpt.graphrag.graphrag_available", lambda: True
        )

        def _boom(*a, **kw):
            raise RuntimeError("extraction blew up")

        monkeypatch.setattr(
            "docsgpt.graphrag.extraction.extract_graph_for_source", _boom
        )
        events = []
        monkeypatch.setattr(
            worker, "publish_user_event",
            lambda user, etype, payload, **kw: events.append((etype, payload)),
        )

        with pytest.raises(RuntimeError):
            worker.extract_graph_worker(task_self, source_id, "alice")

        assert "graph.extract.failed" in [e[0] for e in events]