mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-15 02:16:08 +00:00
Pg 4 (#2390)
* feat: postgres tests * feat: mongo cutoff * feat: mongo cutoff * feat: adjust docs and compose files * fix: mini code mongo removals * fix: tests and k8s mongo stuff * feat: test fixes * fix: ruff * fix: vale * Potential fix for pull request finding 'CodeQL / Clear-text logging of sensitive information' Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> * fix: mini suggestions * vale lint fix 2 * fix: codeql columns thing * fix: test mongo * fix: tests coverage * feat: better tests 4 * feat: more tests * feat: decent coverage * fix: ruff fixes * fix: remove mongo mock * feat: enhance workflow engine and API routes; add document retrieval and source handling * feat: e2e tests * fix: mcp, mongo and more * fix: mini codeql warning * fix: agent chunk view * fix: mini issues * fix: more pg fixes * feat: postgres prep on start * feat: qa tests * fix: mini improvements * fix: tests --------- Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> Co-authored-by: Siddhant Rai <siddhant.rai.5686@gmail.com>
This commit is contained in:
312 files changed
+763725
-25159
No files matched your search
Whitespace-only changes.
@@ -0,0 +1,46 @@
|
||||
"""Fixtures for Celery worker smoke tests.
|
||||
|
||||
These tests exercise the task *bodies* in ``application.worker`` against a
|
||||
real Postgres schema (via the ephemeral ``pg_conn`` fixture from the root
|
||||
``tests/conftest.py``). External I/O — storage, the embedding pipeline,
|
||||
the retriever, the LLM, the backend HTTP callback — is mocked, but every
|
||||
PG write is allowed to hit the real ephemeral DB so the assertions can
|
||||
read the resulting rows back.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import contextmanager
|
||||
from typing import Iterator
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import Connection
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def patch_worker_db(pg_conn, monkeypatch):
|
||||
"""Redirect ``db_session`` / ``db_readonly`` in ``application.worker``.
|
||||
|
||||
Both helpers yield the per-test transactional ``pg_conn``, so any
|
||||
writes a task performs are visible to the test and roll back on
|
||||
teardown. Without this patch the worker would open its own pooled
|
||||
engine and punch past the per-test transaction.
|
||||
"""
|
||||
|
||||
@contextmanager
|
||||
def _use_pg_conn() -> Iterator[Connection]:
|
||||
yield pg_conn
|
||||
|
||||
monkeypatch.setattr("application.worker.db_session", _use_pg_conn)
|
||||
monkeypatch.setattr("application.worker.db_readonly", _use_pg_conn)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def task_self():
|
||||
"""Minimal stand-in for the Celery task ``self`` passed to workers.
|
||||
|
||||
Only ``update_state`` is ever exercised in the happy paths we cover
|
||||
here, so a MagicMock is more than enough.
|
||||
"""
|
||||
return MagicMock(name="celery_task_self")
|
||||
@@ -0,0 +1,159 @@
|
||||
"""Smoke tests for ``agent_webhook_worker`` and ``run_agent_logic``.
|
||||
|
||||
Neither task writes to Postgres directly — they only *read* the agent
|
||||
row (and, in ``run_agent_logic``, the referenced source row). The
|
||||
concrete PG side-effect we assert is therefore a read: the task has
|
||||
to resolve the row from the ephemeral DB to proceed at all; if the
|
||||
lookup returned ``None`` the task would short-circuit with a
|
||||
"not found" error.
|
||||
|
||||
The LLM, retriever, and agent factory are all stubbed — we only care
|
||||
that the PG read path wires up correctly.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from application.storage.db.repositories.agents import AgentsRepository
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestAgentWebhookWorker:
|
||||
def test_resolves_agent_by_uuid_and_runs_logic(
|
||||
self, pg_conn, patch_worker_db, task_self, monkeypatch
|
||||
):
|
||||
from application import worker
|
||||
|
||||
agent = AgentsRepository(pg_conn).create(
|
||||
user_id="alice",
|
||||
name="hook-agent",
|
||||
status="active",
|
||||
agent_type="classic",
|
||||
retriever="classic",
|
||||
chunks=2,
|
||||
key="sk-test-123",
|
||||
)
|
||||
agent_id = str(agent["id"])
|
||||
|
||||
# Capture the resolved agent_config + input; return a fake result.
|
||||
captured: dict = {}
|
||||
|
||||
def _fake_run_agent_logic(agent_config, input_data):
|
||||
captured["config"] = agent_config
|
||||
captured["input"] = input_data
|
||||
return {"answer": "ok", "sources": [], "tool_calls": [], "thought": ""}
|
||||
|
||||
monkeypatch.setattr(worker, "run_agent_logic", _fake_run_agent_logic)
|
||||
|
||||
result = worker.agent_webhook_worker(
|
||||
task_self, agent_id, {"event": "ping"}
|
||||
)
|
||||
|
||||
assert result == {
|
||||
"status": "success",
|
||||
"result": {"answer": "ok", "sources": [], "tool_calls": [], "thought": ""},
|
||||
}
|
||||
# The row pulled from PG is the one we seeded.
|
||||
assert captured["config"]["name"] == "hook-agent"
|
||||
assert str(captured["config"]["id"]) == agent_id
|
||||
assert captured["input"] == '{"event": "ping"}'
|
||||
|
||||
def test_missing_agent_returns_error(
|
||||
self, pg_conn, patch_worker_db, task_self, monkeypatch
|
||||
):
|
||||
from application import worker
|
||||
|
||||
# Run with an id that exists in no shape the resolver understands;
|
||||
# ``looks_like_uuid`` is False so the UUID branch is skipped and
|
||||
# the legacy lookup returns None. Task must yield a clean error
|
||||
# result rather than raising.
|
||||
monkeypatch.setattr(worker, "run_agent_logic", lambda *a, **k: {})
|
||||
result = worker.agent_webhook_worker(task_self, "no-such-agent", {})
|
||||
assert result["status"] == "error"
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestRunAgentLogic:
|
||||
def test_reads_source_row_from_pg(
|
||||
self, pg_conn, patch_worker_db, monkeypatch
|
||||
):
|
||||
"""``run_agent_logic`` looks up the agent's ``source_id`` in PG to
|
||||
pick up the source's ``retriever`` override. Proving the read
|
||||
wires up end-to-end is enough for a smoke test."""
|
||||
from application import worker
|
||||
from application.storage.db.repositories.sources import SourcesRepository
|
||||
|
||||
src = SourcesRepository(pg_conn).create(
|
||||
"src",
|
||||
user_id="alice",
|
||||
type="local",
|
||||
retriever="hybrid",
|
||||
)
|
||||
source_id = str(src["id"])
|
||||
|
||||
# Silence model/provider resolution so we don't need a real key.
|
||||
monkeypatch.setattr(
|
||||
"application.core.model_utils.get_default_model_id", lambda: "gpt-4"
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"application.core.model_utils.validate_model_id", lambda m: True
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"application.core.model_utils.get_provider_from_model_id",
|
||||
lambda m: "openai",
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"application.core.model_utils.get_api_key_for_provider",
|
||||
lambda p: "sk-test",
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"application.utils.calculate_doc_token_budget",
|
||||
lambda model_id=None: 1000,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"application.api.answer.services.stream_processor.get_prompt",
|
||||
lambda prompt_id: "prompt text",
|
||||
)
|
||||
|
||||
# Retriever search returns no docs; agent gen yields a single answer
|
||||
# line so the aggregation loop runs through.
|
||||
captured_source: dict = {}
|
||||
|
||||
class _FakeRetriever:
|
||||
def __init__(self, *args, **kwargs):
|
||||
captured_source.update(kwargs.get("source", {}))
|
||||
|
||||
def search(self, query):
|
||||
return []
|
||||
|
||||
monkeypatch.setattr(
|
||||
"application.retriever.retriever_creator.RetrieverCreator.create_retriever",
|
||||
lambda *a, **kw: _FakeRetriever(**kw),
|
||||
)
|
||||
|
||||
fake_agent = MagicMock(name="agent")
|
||||
fake_agent.gen.return_value = iter([{"answer": "done"}])
|
||||
monkeypatch.setattr(
|
||||
"application.agents.agent_creator.AgentCreator.create_agent",
|
||||
lambda *a, **kw: fake_agent,
|
||||
)
|
||||
|
||||
agent_config = {
|
||||
"id": "agent-uuid",
|
||||
"source_id": source_id,
|
||||
"user_id": "alice",
|
||||
"key": "sk-user",
|
||||
"agent_type": "classic",
|
||||
"chunks": 2,
|
||||
"prompt_id": "default",
|
||||
}
|
||||
|
||||
result = worker.run_agent_logic(agent_config, "hello")
|
||||
|
||||
assert result["answer"] == "done"
|
||||
# Proves the PG read hit the seeded source and its id flowed into
|
||||
# the retriever's ``source={"active_docs": ...}`` param.
|
||||
assert captured_source.get("active_docs") == source_id
|
||||
@@ -0,0 +1,67 @@
|
||||
"""Smoke test for ``application.worker.attachment_worker``.
|
||||
|
||||
The happy path parses an uploaded file and inserts a row into
|
||||
``attachments``. We mock the parser boundary (``StorageCreator.get_storage``
|
||||
returns a storage whose ``process_file`` produces a pre-built Document)
|
||||
but let the PG insert run against the ephemeral ``pg_conn`` so we can
|
||||
assert one concrete row is visible after the task returns.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from application.parser.schema.base import Document
|
||||
from application.storage.db.repositories.attachments import AttachmentsRepository
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestAttachmentWorker:
|
||||
def test_inserts_row_in_attachments(
|
||||
self, pg_conn, patch_worker_db, task_self, monkeypatch
|
||||
):
|
||||
from application import worker
|
||||
|
||||
fake_doc = Document(
|
||||
text="hello world",
|
||||
extra_info={"transcript_language": "en"},
|
||||
)
|
||||
|
||||
fake_storage = MagicMock(name="storage")
|
||||
fake_storage.process_file.return_value = fake_doc
|
||||
monkeypatch.setattr(
|
||||
worker.StorageCreator, "get_storage", lambda: fake_storage
|
||||
)
|
||||
|
||||
# Stub the parser selection so the docling import path isn't taken.
|
||||
monkeypatch.setattr(
|
||||
worker, "get_default_file_extractor", lambda ocr_enabled=False: {}
|
||||
)
|
||||
|
||||
file_info = {
|
||||
"filename": "notes.txt",
|
||||
"attachment_id": "507f1f77bcf86cd799439011",
|
||||
"path": "uploads/user1/notes.txt",
|
||||
"metadata": {"source": "chat"},
|
||||
}
|
||||
|
||||
result = worker.attachment_worker(task_self, file_info, "user1")
|
||||
|
||||
assert result["filename"] == "notes.txt"
|
||||
assert result["token_count"] > 0
|
||||
# Parser metadata (``transcript_*``) should have been merged in.
|
||||
assert result["metadata"]["transcript_language"] == "en"
|
||||
assert result["metadata"]["source"] == "chat"
|
||||
|
||||
# Row should be resolvable by the caller-visible handle stored in
|
||||
# ``legacy_mongo_id``.
|
||||
row = AttachmentsRepository(pg_conn).get_by_legacy_id(
|
||||
file_info["attachment_id"], "user1"
|
||||
)
|
||||
assert row is not None, "attachment_worker should insert a row"
|
||||
assert row["filename"] == "notes.txt"
|
||||
assert row["upload_path"] == "uploads/user1/notes.txt"
|
||||
assert row["content"] == "hello world"
|
||||
assert row["user_id"] == "user1"
|
||||
@@ -0,0 +1,109 @@
|
||||
"""Smoke test for ``application.worker.ingest_connector`` in sync mode.
|
||||
|
||||
Sync mode (``operation_mode="sync"``) bumps ``sources.date`` on the
|
||||
target source row, the same PG side-effect as ``remote_worker``. Upload
|
||||
mode does not write to PG directly — that happens indirectly via the
|
||||
``upload_index`` HTTP callback to the backend — so the happy-path smoke
|
||||
is the sync variant.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from application.parser.schema.base import Document
|
||||
from application.storage.db.repositories.sources import SourcesRepository
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def _mock_connector_pipeline(monkeypatch):
|
||||
"""Stub the connector + pipeline so only PG writes are real."""
|
||||
from application import worker
|
||||
|
||||
fake_connector = MagicMock(name="connector")
|
||||
fake_connector.download_to_directory.return_value = {
|
||||
"files_downloaded": 1,
|
||||
"empty_result": False,
|
||||
}
|
||||
monkeypatch.setattr(
|
||||
worker.ConnectorCreator, "is_supported", staticmethod(lambda s: True)
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
worker.ConnectorCreator,
|
||||
"create_connector",
|
||||
staticmethod(lambda source_type, session_token: fake_connector),
|
||||
)
|
||||
|
||||
fake_reader = MagicMock(name="reader")
|
||||
fake_reader.load_data.return_value = [
|
||||
Document(
|
||||
text="connector body",
|
||||
extra_info={"source": "connector/file.md", "file_path": "file.md"},
|
||||
)
|
||||
]
|
||||
fake_reader.directory_structure = {
|
||||
"file.md": {
|
||||
"type": "text/markdown",
|
||||
"size_bytes": 12,
|
||||
"token_count": 3,
|
||||
}
|
||||
}
|
||||
monkeypatch.setattr(
|
||||
worker, "SimpleDirectoryReader", lambda *a, **kw: fake_reader
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
worker,
|
||||
"embed_and_store_documents",
|
||||
lambda docs, full_path, source_id, task: None,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
worker, "upload_index", lambda full_path, file_data: None
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestIngestConnectorSyncUpdatesDate:
|
||||
def test_sync_mode_bumps_source_date(
|
||||
self,
|
||||
pg_conn,
|
||||
patch_worker_db,
|
||||
task_self,
|
||||
monkeypatch,
|
||||
_mock_connector_pipeline,
|
||||
):
|
||||
from application import worker
|
||||
import datetime as dt
|
||||
|
||||
old_date = dt.datetime(2019, 6, 1, tzinfo=dt.timezone.utc)
|
||||
src = SourcesRepository(pg_conn).create(
|
||||
"gdrive-folder",
|
||||
user_id="dave",
|
||||
type="connector:file",
|
||||
retriever="classic",
|
||||
sync_frequency="weekly",
|
||||
remote_data={"provider": "google_drive"},
|
||||
date=old_date,
|
||||
)
|
||||
source_id = str(src["id"])
|
||||
|
||||
worker.ingest_connector(
|
||||
task_self,
|
||||
"gdrive-folder",
|
||||
"dave",
|
||||
"google_drive",
|
||||
session_token="tok",
|
||||
file_ids=["f1"],
|
||||
folder_ids=[],
|
||||
operation_mode="sync",
|
||||
doc_id=source_id,
|
||||
sync_frequency="weekly",
|
||||
)
|
||||
|
||||
refreshed = SourcesRepository(pg_conn).get(source_id, "dave")
|
||||
assert refreshed is not None
|
||||
assert refreshed["date"] > old_date, (
|
||||
"ingest_connector(sync) should push sources.date forward"
|
||||
)
|
||||
@@ -0,0 +1,93 @@
|
||||
"""Smoke test for ``application.worker.ingest_worker``.
|
||||
|
||||
``ingest_worker`` does **not** write to Postgres directly. Its PG
|
||||
side-effect (creating the ``sources`` row) goes through the
|
||||
``upload_index`` HTTP callback to the backend, which writes to PG in
|
||||
its own request context. That callback is intentionally out of scope
|
||||
here — we can't reach it from the worker test without spinning up the
|
||||
Flask app.
|
||||
|
||||
What we *can* smoke: the task body runs end-to-end, the pipeline gets
|
||||
invoked with the expected job metadata, and ``upload_index`` is handed
|
||||
a ``file_data`` payload that carries the caller-provided ``user``,
|
||||
``job_name``, and ``retriever``. That's the contract the
|
||||
backend-facing write depends on.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from application.parser.schema.base import Document
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestIngestWorker:
|
||||
def test_invokes_upload_index_with_expected_payload(
|
||||
self, patch_worker_db, task_self, monkeypatch
|
||||
):
|
||||
from application import worker
|
||||
|
||||
# Storage: pretend a single file was uploaded at ``inputs/.../a.txt``.
|
||||
from io import BytesIO
|
||||
fake_storage = MagicMock(name="storage")
|
||||
fake_storage.is_directory.return_value = False
|
||||
fake_storage.get_file.return_value = BytesIO(b"hello")
|
||||
monkeypatch.setattr(
|
||||
worker.StorageCreator, "get_storage", lambda: fake_storage
|
||||
)
|
||||
|
||||
# Reader: return a parsed doc.
|
||||
fake_reader = MagicMock(name="reader")
|
||||
fake_reader.load_data.return_value = [
|
||||
Document(text="hello body", extra_info={"source": "a.txt"})
|
||||
]
|
||||
fake_reader.directory_structure = {
|
||||
"a.txt": {
|
||||
"type": "text/plain",
|
||||
"size_bytes": 5,
|
||||
"token_count": 2,
|
||||
}
|
||||
}
|
||||
monkeypatch.setattr(
|
||||
worker, "SimpleDirectoryReader", lambda *a, **kw: fake_reader
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
worker,
|
||||
"embed_and_store_documents",
|
||||
lambda docs, full_path, source_id, task: None,
|
||||
)
|
||||
|
||||
captured: list[dict] = []
|
||||
|
||||
def _fake_upload(full_path, file_data):
|
||||
captured.append(file_data)
|
||||
|
||||
monkeypatch.setattr(worker, "upload_index", _fake_upload)
|
||||
|
||||
result = worker.ingest_worker(
|
||||
task_self,
|
||||
directory="inputs",
|
||||
formats=[".txt"],
|
||||
job_name="job1",
|
||||
file_path="inputs/eve/job1/a.txt",
|
||||
filename="a.txt",
|
||||
user="eve",
|
||||
retriever="classic",
|
||||
)
|
||||
|
||||
assert result["limited"] is False
|
||||
assert result["user"] == "eve"
|
||||
assert result["name_job"] == "job1"
|
||||
|
||||
assert len(captured) == 1
|
||||
payload = captured[0]
|
||||
assert payload["user"] == "eve"
|
||||
assert payload["name"] == "job1"
|
||||
assert payload["retriever"] == "classic"
|
||||
assert payload["type"] == "local"
|
||||
# A fresh source UUID is minted for the backend /upload_index route.
|
||||
assert payload["id"]
|
||||
@@ -0,0 +1,93 @@
|
||||
"""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 unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
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"]
|
||||
@@ -0,0 +1,212 @@
|
||||
"""Smoke tests for ``application.worker.remote_worker`` and ``sync_worker``.
|
||||
|
||||
``remote_worker`` in ``sync`` mode does one PG write: it bumps
|
||||
``sources.date`` on the referenced source row to ``now()``. That's the
|
||||
side-effect we assert here. The remote loader, chunker, embedding
|
||||
pipeline, and the backend HTTP callback are all mocked — only the PG
|
||||
update is real.
|
||||
|
||||
``sync_worker`` reads rows out of ``sources`` whose ``sync_frequency``
|
||||
matches and dispatches them through ``sync`` → ``remote_worker``. We
|
||||
assert one seeded row is discovered and forwarded.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from application.parser.schema.base import Document
|
||||
from application.storage.db.repositories.sources import SourcesRepository
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def _mock_remote_pipeline(monkeypatch):
|
||||
"""Stub out the non-PG boundaries used by ``remote_worker``."""
|
||||
from application import worker
|
||||
|
||||
fake_loader = MagicMock(name="remote_loader")
|
||||
fake_loader.load_data.return_value = [
|
||||
Document(
|
||||
text="page body",
|
||||
extra_info={"file_path": "guides/setup.md", "title": "setup"},
|
||||
doc_id="d1",
|
||||
)
|
||||
]
|
||||
monkeypatch.setattr(
|
||||
worker.RemoteCreator, "create_loader", lambda loader: fake_loader
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
worker,
|
||||
"embed_and_store_documents",
|
||||
lambda docs, full_path, source_id, task: None,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
worker, "upload_index", lambda full_path, file_data: None
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestRemoteWorkerSyncUpdatesDate:
|
||||
def test_sync_mode_bumps_source_date(
|
||||
self,
|
||||
pg_conn,
|
||||
patch_worker_db,
|
||||
task_self,
|
||||
monkeypatch,
|
||||
_mock_remote_pipeline,
|
||||
):
|
||||
from application import worker
|
||||
|
||||
# Seed a source with a known old ``date`` we can compare against.
|
||||
import datetime as dt
|
||||
|
||||
old_date = dt.datetime(2020, 1, 1, tzinfo=dt.timezone.utc)
|
||||
src = SourcesRepository(pg_conn).create(
|
||||
"my-remote",
|
||||
user_id="bob",
|
||||
type="crawler",
|
||||
retriever="classic",
|
||||
sync_frequency="daily",
|
||||
remote_data={"urls": ["http://example.com"]},
|
||||
date=old_date,
|
||||
)
|
||||
source_id = str(src["id"])
|
||||
|
||||
worker.remote_worker(
|
||||
task_self,
|
||||
{"urls": ["http://example.com"]},
|
||||
"my-remote",
|
||||
"bob",
|
||||
"crawler",
|
||||
directory="temp",
|
||||
retriever="classic",
|
||||
sync_frequency="daily",
|
||||
operation_mode="sync",
|
||||
doc_id=source_id,
|
||||
)
|
||||
|
||||
refreshed = SourcesRepository(pg_conn).get(source_id, "bob")
|
||||
assert refreshed is not None
|
||||
assert refreshed["date"] > old_date, (
|
||||
"remote_worker(sync) should push sources.date forward"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestSyncWorker:
|
||||
def test_reads_sources_and_dispatches_sync(
|
||||
self,
|
||||
pg_conn,
|
||||
patch_worker_db,
|
||||
task_self,
|
||||
monkeypatch,
|
||||
):
|
||||
"""``sync_worker`` selects rows by ``sync_frequency`` and passes
|
||||
each to ``sync``. We assert the seeded row is discovered and
|
||||
forwarded with the right doc_id — the nested ``sync`` call is
|
||||
stubbed so we don't re-run the whole remote pipeline here."""
|
||||
from application import worker
|
||||
|
||||
src = SourcesRepository(pg_conn).create(
|
||||
"weekly-feed",
|
||||
user_id="carol",
|
||||
type="url",
|
||||
retriever="classic",
|
||||
sync_frequency="weekly",
|
||||
remote_data={"url": "http://example.com"},
|
||||
)
|
||||
|
||||
captured: list[dict] = []
|
||||
|
||||
def _fake_sync(self, source_data, name_job, user, loader,
|
||||
sync_frequency, retriever, doc_id=None, directory="temp"):
|
||||
captured.append({
|
||||
"name_job": name_job,
|
||||
"user": user,
|
||||
"loader": loader,
|
||||
"sync_frequency": sync_frequency,
|
||||
"retriever": retriever,
|
||||
"doc_id": doc_id,
|
||||
})
|
||||
return {"status": "success"}
|
||||
|
||||
monkeypatch.setattr(worker, "sync", _fake_sync)
|
||||
|
||||
result = worker.sync_worker(task_self, "weekly")
|
||||
|
||||
assert result["total_sync_count"] == 1
|
||||
assert result["sync_success"] == 1
|
||||
assert len(captured) == 1
|
||||
assert captured[0]["name_job"] == "weekly-feed"
|
||||
assert captured[0]["user"] == "carol"
|
||||
assert captured[0]["loader"] == "url"
|
||||
assert captured[0]["doc_id"] == str(src["id"])
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestRemoteWorkerPathTraversal:
|
||||
"""Regression: ``name_job`` must not be usable as a path segment.
|
||||
|
||||
Historically ``remote_worker`` built its workspace from
|
||||
``os.path.join(directory, user, name_job)`` and cleaned it up with
|
||||
``shutil.rmtree`` in a ``finally``. A ``name_job`` like
|
||||
``../../evil`` would therefore let an authenticated caller delete
|
||||
directories outside the intended ``<directory>/<user>/`` root.
|
||||
The fix uses a random uuid leaf; ``name_job`` is metadata only.
|
||||
"""
|
||||
|
||||
def test_traversal_name_job_does_not_escape_user_workspace(
|
||||
self,
|
||||
tmp_path,
|
||||
task_self,
|
||||
monkeypatch,
|
||||
_mock_remote_pipeline,
|
||||
):
|
||||
from application import worker
|
||||
|
||||
created_paths: list[str] = []
|
||||
deleted_paths: list[str] = []
|
||||
real_makedirs = os.makedirs
|
||||
real_rmtree = worker.shutil.rmtree
|
||||
|
||||
def _spy_makedirs(path, *args, **kwargs):
|
||||
created_paths.append(path)
|
||||
return real_makedirs(path, *args, **kwargs)
|
||||
|
||||
def _spy_rmtree(path, *args, **kwargs):
|
||||
deleted_paths.append(path)
|
||||
return real_rmtree(path, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(worker.os, "makedirs", _spy_makedirs)
|
||||
monkeypatch.setattr(worker.shutil, "rmtree", _spy_rmtree)
|
||||
|
||||
directory = str(tmp_path / "temp")
|
||||
user = "bob"
|
||||
malicious_name = "../../evil"
|
||||
|
||||
worker.remote_worker(
|
||||
task_self,
|
||||
{"urls": ["http://example.com"]},
|
||||
malicious_name,
|
||||
user,
|
||||
"crawler",
|
||||
directory=directory,
|
||||
operation_mode="upload",
|
||||
)
|
||||
|
||||
directory_real = os.path.realpath(directory)
|
||||
user_root = os.path.realpath(os.path.join(directory, user))
|
||||
|
||||
rmtree_targets = [
|
||||
os.path.realpath(p)
|
||||
for p in deleted_paths
|
||||
if os.path.realpath(p).startswith(directory_real)
|
||||
]
|
||||
assert len(rmtree_targets) == 1, rmtree_targets
|
||||
assert rmtree_targets[0].startswith(user_root + os.sep), (
|
||||
f"rmtree target {rmtree_targets[0]} escaped {user_root}"
|
||||
)
|
||||
assert malicious_name not in "".join(created_paths + deleted_paths)
|
||||
Reference in new issue
Block a user