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
+138
-123
@@ -1,108 +1,125 @@
|
||||
"""Unit tests for TodoListTool.
|
||||
|
||||
These tests replace the old Mongo-coupled suite. They work by:
|
||||
|
||||
* Patching ``TodosRepository`` inside the tool module with a small
|
||||
in-memory fake that implements only the methods the tool calls.
|
||||
* Patching ``db_session`` / ``db_readonly`` inside the tool module with
|
||||
a no-op context manager so the tool never tries to open a real
|
||||
connection.
|
||||
* Exercising the tool with a real UUID ``tool_id`` so ``_pg_enabled()``
|
||||
returns True, plus a sentinel ``default_{uid}`` case so the short-
|
||||
circuit path is covered too.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from contextlib import contextmanager
|
||||
|
||||
import pytest
|
||||
|
||||
from application.agents.tools.todo_list import TodoListTool
|
||||
from application.core.settings import settings
|
||||
|
||||
|
||||
class FakeCursor(list):
|
||||
def sort(self, key, direction):
|
||||
reverse = direction == -1
|
||||
sorted_list = sorted(self, key=lambda d: d.get(key, 0), reverse=reverse)
|
||||
return FakeCursor(sorted_list)
|
||||
class _FakeTodosRepo:
|
||||
"""In-memory stand-in for ``TodosRepository``.
|
||||
|
||||
def limit(self, count):
|
||||
return FakeCursor(self[:count])
|
||||
Implements only the methods called by ``TodoListTool``. Shared state
|
||||
lives on the class so multiple instances (one per ``with`` block)
|
||||
see the same data, mirroring a real DB.
|
||||
"""
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
_store: list[dict] = []
|
||||
_next_ids: dict[str, int] = {}
|
||||
|
||||
def __next__(self):
|
||||
if not self:
|
||||
raise StopIteration
|
||||
return self.pop(0)
|
||||
def __init__(self, conn=None) -> None:
|
||||
self._conn = conn
|
||||
|
||||
@classmethod
|
||||
def reset(cls) -> None:
|
||||
cls._store = []
|
||||
cls._next_ids = {}
|
||||
|
||||
class FakeCollection:
|
||||
def __init__(self):
|
||||
self.docs = {}
|
||||
self._id_counter = 0
|
||||
def create(self, user_id, tool_id, title, *, todo_id=None, legacy_mongo_id=None):
|
||||
if todo_id is None:
|
||||
todo_id = self._next_ids.get(tool_id, 0) + 1
|
||||
self._next_ids[tool_id] = max(self._next_ids.get(tool_id, 0), todo_id)
|
||||
row = {
|
||||
"id": str(uuid.uuid4()),
|
||||
"user_id": user_id,
|
||||
"tool_id": tool_id,
|
||||
"todo_id": todo_id,
|
||||
"title": title,
|
||||
"completed": False,
|
||||
"legacy_mongo_id": legacy_mongo_id,
|
||||
}
|
||||
self._store.append(row)
|
||||
return row
|
||||
|
||||
def _generate_id(self):
|
||||
self._id_counter += 1
|
||||
return f"fake_id_{self._id_counter}"
|
||||
|
||||
def create_index(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
def insert_one(self, doc):
|
||||
key = (doc["user_id"], doc["tool_id"], doc["todo_id"])
|
||||
if "_id" not in doc:
|
||||
doc["_id"] = self._generate_id()
|
||||
self.docs[key] = doc
|
||||
return type("res", (), {"inserted_id": doc["_id"]})
|
||||
|
||||
def find_one(self, query):
|
||||
key = (query.get("user_id"), query.get("tool_id"), query.get("todo_id"))
|
||||
return self.docs.get(key)
|
||||
|
||||
def find(self, query, projection=None):
|
||||
user_id = query.get("user_id")
|
||||
tool_id = query.get("tool_id")
|
||||
filtered = [
|
||||
doc for (uid, tid, _), doc in self.docs.items()
|
||||
if uid == user_id and tid == tool_id
|
||||
def list_for_tool(self, user_id, tool_id):
|
||||
return [
|
||||
r for r in self._store
|
||||
if r["user_id"] == user_id and r["tool_id"] == tool_id
|
||||
]
|
||||
return FakeCursor(filtered)
|
||||
|
||||
def update_one(self, query, update, upsert=False):
|
||||
key = (query.get("user_id"), query.get("tool_id"), query.get("todo_id"))
|
||||
if key in self.docs:
|
||||
self.docs[key].update(update.get("$set", {}))
|
||||
return type("res", (), {"matched_count": 1})
|
||||
elif upsert:
|
||||
new_doc = {**query, **update.get("$set", {}), "_id": self._generate_id()}
|
||||
self.docs[key] = new_doc
|
||||
return type("res", (), {"matched_count": 1})
|
||||
else:
|
||||
return type("res", (), {"matched_count": 0})
|
||||
|
||||
def find_one_and_update(self, query, update):
|
||||
key = (query.get("user_id"), query.get("tool_id"), query.get("todo_id"))
|
||||
if key in self.docs:
|
||||
self.docs[key].update(update.get("$set", {}))
|
||||
return self.docs[key]
|
||||
def get_by_tool_and_todo_id(self, user_id, tool_id, todo_id):
|
||||
for r in self._store:
|
||||
if (r["user_id"] == user_id
|
||||
and r["tool_id"] == tool_id
|
||||
and r["todo_id"] == todo_id):
|
||||
return r
|
||||
return None
|
||||
|
||||
def find_one_and_delete(self, query):
|
||||
key = (query.get("user_id"), query.get("tool_id"), query.get("todo_id"))
|
||||
if key in self.docs:
|
||||
return self.docs.pop(key)
|
||||
return None
|
||||
def update_title_by_tool_and_todo_id(self, user_id, tool_id, todo_id, title):
|
||||
row = self.get_by_tool_and_todo_id(user_id, tool_id, todo_id)
|
||||
if row is None:
|
||||
return False
|
||||
row["title"] = title
|
||||
return True
|
||||
|
||||
def delete_one(self, query):
|
||||
key = (query.get("user_id"), query.get("tool_id"), query.get("todo_id"))
|
||||
if key in self.docs:
|
||||
del self.docs[key]
|
||||
return type("res", (), {"deleted_count": 1})
|
||||
return type("res", (), {"deleted_count": 0})
|
||||
def set_completed(self, user_id, tool_id, todo_id, completed):
|
||||
row = self.get_by_tool_and_todo_id(user_id, tool_id, todo_id)
|
||||
if row is None:
|
||||
return False
|
||||
row["completed"] = bool(completed)
|
||||
return True
|
||||
|
||||
def delete_by_tool_and_todo_id(self, user_id, tool_id, todo_id):
|
||||
for idx, r in enumerate(self._store):
|
||||
if (r["user_id"] == user_id
|
||||
and r["tool_id"] == tool_id
|
||||
and r["todo_id"] == todo_id):
|
||||
self._store.pop(idx)
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _noop_conn():
|
||||
yield None
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def todo_tool(monkeypatch) -> TodoListTool:
|
||||
"""Provides a TodoListTool with a fake MongoDB backend."""
|
||||
# Reset the MongoDB client cache to ensure our mock is used
|
||||
from application.core.mongo_db import MongoDB
|
||||
MongoDB._client = None
|
||||
|
||||
fake_collection = FakeCollection()
|
||||
fake_client = {settings.MONGO_DB_NAME: {"todos": fake_collection}}
|
||||
monkeypatch.setattr("application.core.mongo_db.MongoDB.get_client", lambda: fake_client)
|
||||
return TodoListTool({"tool_id": "test_tool"}, user_id="test_user")
|
||||
def todo_tool(monkeypatch):
|
||||
"""Return a ``TodoListTool`` wired to the in-memory fake repo."""
|
||||
_FakeTodosRepo.reset()
|
||||
monkeypatch.setattr(
|
||||
"application.agents.tools.todo_list.TodosRepository", _FakeTodosRepo
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"application.agents.tools.todo_list.db_session", _noop_conn
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"application.agents.tools.todo_list.db_readonly", _noop_conn
|
||||
)
|
||||
# Real UUID so ``_pg_enabled()`` returns True.
|
||||
return TodoListTool({"tool_id": str(uuid.uuid4())}, user_id="test_user")
|
||||
|
||||
|
||||
def test_create_and_get(todo_tool: TodoListTool):
|
||||
def test_create_and_get(todo_tool):
|
||||
res = todo_tool.execute_action("create", title="Write tests")
|
||||
assert "Todo created with ID" in res
|
||||
# Extract todo_id from response like "Todo created with ID test_user_test_tool_1: Write tests"
|
||||
todo_id = res.split("ID ")[1].split(":")[0].strip()
|
||||
|
||||
get_res = todo_tool.execute_action("get", todo_id=todo_id)
|
||||
@@ -110,7 +127,7 @@ def test_create_and_get(todo_tool: TodoListTool):
|
||||
assert "Write tests" in get_res
|
||||
|
||||
|
||||
def test_get_all_todos(todo_tool: TodoListTool):
|
||||
def test_get_all_todos(todo_tool):
|
||||
todo_tool.execute_action("create", title="Task 1")
|
||||
todo_tool.execute_action("create", title="Task 2")
|
||||
|
||||
@@ -119,7 +136,7 @@ def test_get_all_todos(todo_tool: TodoListTool):
|
||||
assert "Task 2" in list_res
|
||||
|
||||
|
||||
def test_update_todo(todo_tool: TodoListTool):
|
||||
def test_update_todo(todo_tool):
|
||||
create_res = todo_tool.execute_action("create", title="Initial Title")
|
||||
todo_id = create_res.split("ID ")[1].split(":")[0].strip()
|
||||
|
||||
@@ -131,24 +148,21 @@ def test_update_todo(todo_tool: TodoListTool):
|
||||
assert "Updated Title" in get_res
|
||||
|
||||
|
||||
def test_complete_todo(todo_tool: TodoListTool):
|
||||
def test_complete_todo(todo_tool):
|
||||
create_res = todo_tool.execute_action("create", title="To Complete")
|
||||
todo_id = create_res.split("ID ")[1].split(":")[0].strip()
|
||||
|
||||
# Check initial status is open
|
||||
get_res = todo_tool.execute_action("get", todo_id=todo_id)
|
||||
assert "open" in get_res
|
||||
|
||||
# Mark as completed
|
||||
complete_res = todo_tool.execute_action("complete", todo_id=todo_id)
|
||||
assert "completed" in complete_res.lower()
|
||||
|
||||
# Verify status changed to completed
|
||||
get_res = todo_tool.execute_action("get", todo_id=todo_id)
|
||||
assert "completed" in get_res
|
||||
|
||||
|
||||
def test_delete_todo(todo_tool: TodoListTool):
|
||||
def test_delete_todo(todo_tool):
|
||||
create_res = todo_tool.execute_action("create", title="To Delete")
|
||||
todo_id = create_res.split("ID ")[1].split(":")[0].strip()
|
||||
|
||||
@@ -160,45 +174,46 @@ def test_delete_todo(todo_tool: TodoListTool):
|
||||
assert "not found" in get_res
|
||||
|
||||
|
||||
def test_isolation_and_default_tool_id(monkeypatch):
|
||||
"""Ensure todos are isolated by tool_id and user_id."""
|
||||
# Reset the MongoDB client cache to ensure our mock is used
|
||||
from application.core.mongo_db import MongoDB
|
||||
MongoDB._client = None
|
||||
def test_isolation_per_tool_id(monkeypatch):
|
||||
"""Todos created under one tool_id are not visible to another tool_id."""
|
||||
_FakeTodosRepo.reset()
|
||||
monkeypatch.setattr(
|
||||
"application.agents.tools.todo_list.TodosRepository", _FakeTodosRepo
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"application.agents.tools.todo_list.db_session", _noop_conn
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"application.agents.tools.todo_list.db_readonly", _noop_conn
|
||||
)
|
||||
|
||||
fake_collection = FakeCollection()
|
||||
fake_client = {settings.MONGO_DB_NAME: {"todos": fake_collection}}
|
||||
monkeypatch.setattr("application.core.mongo_db.MongoDB.get_client", lambda: fake_client)
|
||||
tool1 = TodoListTool({"tool_id": str(uuid.uuid4())}, user_id="u1")
|
||||
tool2 = TodoListTool({"tool_id": str(uuid.uuid4())}, user_id="u1")
|
||||
|
||||
# Same user, different tool_id
|
||||
tool1 = TodoListTool({"tool_id": "tool_1"}, user_id="u1")
|
||||
tool2 = TodoListTool({"tool_id": "tool_2"}, user_id="u1")
|
||||
r1 = tool1.execute_action("create", title="from tool 1")
|
||||
tool2.execute_action("create", title="from tool 2")
|
||||
|
||||
r1_create = tool1.execute_action("create", title="from tool 1")
|
||||
r2_create = tool2.execute_action("create", title="from tool 2")
|
||||
todo_id_1 = r1.split("ID ")[1].split(":")[0].strip()
|
||||
|
||||
todo_id_1 = r1_create.split("ID ")[1].split(":")[0].strip()
|
||||
todo_id_2 = r2_create.split("ID ")[1].split(":")[0].strip()
|
||||
# tool2 cannot see tool1's todo even though same user_id + same todo_id=1.
|
||||
res = tool2.execute_action("get", todo_id=todo_id_1)
|
||||
# tool2 has its own todo_id=1 ("from tool 2"), so the right thing to
|
||||
# verify is isolation: tool1's title does not leak into tool2's view.
|
||||
assert "from tool 1" not in res
|
||||
|
||||
r1 = tool1.execute_action("get", todo_id=todo_id_1)
|
||||
r2 = tool2.execute_action("get", todo_id=todo_id_2)
|
||||
|
||||
assert "Error" not in r1
|
||||
assert "from tool 1" in r1
|
||||
def test_sentinel_tool_id_short_circuits():
|
||||
"""A ``default_{user_id}`` tool_id must no-op with a polite error."""
|
||||
tool = TodoListTool({}, user_id="default_user")
|
||||
assert tool.tool_id == "default_default_user"
|
||||
|
||||
assert "Error" not in r2
|
||||
assert "from tool 2" in r2
|
||||
res = tool.execute_action("create", title="shared default")
|
||||
assert "Error" in res
|
||||
assert "not configured" in res.lower() or "unavailable" in res.lower()
|
||||
|
||||
# Same user, no tool_id → should default to same value
|
||||
t3 = TodoListTool({}, user_id="default_user")
|
||||
t4 = TodoListTool({}, user_id="default_user")
|
||||
|
||||
assert t3.tool_id == "default_default_user"
|
||||
assert t4.tool_id == "default_default_user"
|
||||
|
||||
create_res = t3.execute_action("create", title="shared default")
|
||||
todo_id = create_res.split("ID ")[1].split(":")[0].strip()
|
||||
r = t4.execute_action("get", todo_id=todo_id)
|
||||
|
||||
assert "Error" not in r
|
||||
assert "shared default" in r
|
||||
def test_no_user_id_returns_error():
|
||||
tool = TodoListTool({"tool_id": str(uuid.uuid4())}, user_id=None)
|
||||
res = tool.execute_action("list")
|
||||
assert "Error" in res
|
||||
assert "user_id" in res
|
||||
Reference in new issue
Block a user