* 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:
authored and GitHub committed 2026-04-18 13:13:57 +01:00
1 parent 7f6360b4ff
commit 81b6ee5daa
312 files changed
+763725 -25159

No files matched your search

View File
Whitespace-only changes.
+46
View File
@@ -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")
+159
View File
@@ -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
+67
View File
@@ -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"
+109
View File
@@ -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"
)
+93
View File
@@ -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"]
+212
View File
@@ -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)