tiennm99/DocsGPT
1
0 Fork 0
mirror of https://github.com/tiennm99/DocsGPT.git synced 2026-10-18 12:13:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

Pg 4 (#2390)

Browse Source 
* 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:
AlexCopilot Autofix powered by AISiddhant Rai
 authored and GitHub committed 2026-04-18 13:13:57 +01:00
1 parent 7f6360b4ff
commit 81b6ee5daa
312 files changed
+763725 -25159
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/api/answer/routes/test_answer.py
+10 -6
View File
@@ -1,10 +1,14 @@
"""Tests for application/api/answer/routes/answer.py"""
import json
import uuid
from unittest.mock import MagicMock, patch
import pytest
from bson import ObjectId
# Static IDs
_CONV_ID = "507f1f77bcf86cd799439011"
_AGENT_ID = "507f1f77bcf86cd799439012"
@pytest.fixture
@@ -15,9 +19,9 @@ def mock_stream_processor():
) as MockProcessor:
processor = MagicMock()
processor.decoded_token = {"sub": "test_user"}
processor.conversation_id = str(ObjectId())
processor.conversation_id = _CONV_ID
processor.agent_config = {}
processor.agent_id = str(ObjectId())
processor.agent_id = _AGENT_ID
processor.is_shared_usage = False
processor.shared_token = None
processor.model_id = "gpt-4"
@@ -50,7 +54,7 @@ class TestAnswerResourcePost:
assert resp.status_code == 400
def test_successful_answer(self, answer_client, mock_stream_processor):
conv_id = str(ObjectId())
conv_id = str(uuid.uuid4())
with patch.object(
mock_stream_processor.build_agent.return_value,
"gen",
@@ -161,7 +165,7 @@ class TestAnswerResourcePost:
def test_structured_info_merged_into_result(
self, answer_client, mock_stream_processor
):
conv_id = str(ObjectId())
conv_id = str(uuid.uuid4())
with patch(
"application.api.answer.routes.answer.AnswerResource.validate_request",
return_value=None,
@@ -188,7 +192,7 @@ class TestAnswerResourcePost:
def test_result_contains_all_expected_fields(
self, answer_client, mock_stream_processor
):
conv_id = str(ObjectId())
conv_id = str(uuid.uuid4())
with patch(
"application.api.answer.routes.answer.AnswerResource.validate_request",
return_value=None,
tests/api/answer/routes/test_base.py
+143 -237
View File
@@ -1,12 +1,14 @@
import datetime
import uuid
from contextlib import contextmanager
from unittest.mock import MagicMock, patch
import pytest
from bson import ObjectId
@pytest.mark.unit
class TestBaseAnswerValidation:
pass
def test_validate_request_passes_with_required_fields(
self, mock_mongo_db, flask_app
):
@@ -51,7 +53,7 @@ class TestBaseAnswerValidation:
with flask_app.app_context():
resource = BaseAnswerResource()
data = {"question": "Test", "conversation_id": str(ObjectId())}
data = {"question": "Test", "conversation_id": str(uuid.uuid4())}
result = resource.validate_request(data, require_conversation_id=True)
@@ -60,6 +62,8 @@ class TestBaseAnswerValidation:
@pytest.mark.unit
class TestUsageChecking:
pass
def test_returns_none_when_no_api_key(self, mock_mongo_db, flask_app):
from application.api.answer.routes.base import BaseAnswerResource
@@ -71,203 +75,18 @@ class TestUsageChecking:
assert result is None
def test_returns_error_for_invalid_api_key(self, mock_mongo_db, flask_app):
from application.api.answer.routes.base import BaseAnswerResource
with flask_app.app_context():
resource = BaseAnswerResource()
agent_config = {"user_api_key": "invalid_key_123"}
result = resource.check_usage(agent_config)
assert result is not None
assert result.status_code == 401
assert result.json["success"] is False
assert "invalid" in result.json["message"].lower()
def test_checks_token_limit_when_enabled(self, mock_mongo_db, flask_app):
from application.api.answer.routes.base import BaseAnswerResource
from application.core.settings import settings
with flask_app.app_context():
agents_collection = mock_mongo_db[settings.MONGO_DB_NAME]["agents"]
agent_id = ObjectId()
agents_collection.insert_one(
{
"_id": agent_id,
"key": "test_key",
"limited_token_mode": True,
"token_limit": 1000,
"limited_request_mode": False,
}
)
resource = BaseAnswerResource()
agent_config = {"user_api_key": "test_key"}
result = resource.check_usage(agent_config)
assert result is None
def test_checks_request_limit_when_enabled(self, mock_mongo_db, flask_app):
from application.api.answer.routes.base import BaseAnswerResource
from application.core.settings import settings
with flask_app.app_context():
agents_collection = mock_mongo_db[settings.MONGO_DB_NAME]["agents"]
agent_id = ObjectId()
agents_collection.insert_one(
{
"_id": agent_id,
"key": "test_key",
"limited_token_mode": False,
"limited_request_mode": True,
"request_limit": 100,
}
)
resource = BaseAnswerResource()
agent_config = {"user_api_key": "test_key"}
result = resource.check_usage(agent_config)
assert result is None
def test_uses_default_limits_when_not_specified(self, mock_mongo_db, flask_app):
from application.api.answer.routes.base import BaseAnswerResource
from application.core.settings import settings
with flask_app.app_context():
agents_collection = mock_mongo_db[settings.MONGO_DB_NAME]["agents"]
agent_id = ObjectId()
agents_collection.insert_one(
{
"_id": agent_id,
"key": "test_key",
"limited_token_mode": True,
"limited_request_mode": True,
}
)
resource = BaseAnswerResource()
agent_config = {"user_api_key": "test_key"}
result = resource.check_usage(agent_config)
assert result is None
def test_exceeds_token_limit(self, mock_mongo_db, flask_app):
from application.api.answer.routes.base import BaseAnswerResource
from application.core.settings import settings
with flask_app.app_context():
agents_collection = mock_mongo_db[settings.MONGO_DB_NAME]["agents"]
token_usage_collection = mock_mongo_db[settings.MONGO_DB_NAME][
"token_usage"
]
agent_id = ObjectId()
agents_collection.insert_one(
{
"_id": agent_id,
"key": "test_key",
"limited_token_mode": True,
"token_limit": 100,
"limited_request_mode": False,
}
)
token_usage_collection.insert_one(
{
"_id": ObjectId(),
"api_key": "test_key",
"prompt_tokens": 60,
"generated_tokens": 50,
"timestamp": datetime.datetime.now(),
}
)
resource = BaseAnswerResource()
agent_config = {"user_api_key": "test_key"}
result = resource.check_usage(agent_config)
assert result is not None
assert result.status_code == 429
assert result.json["success"] is False
assert "usage limit" in result.json["message"].lower()
def test_exceeds_request_limit(self, mock_mongo_db, flask_app):
from application.api.answer.routes.base import BaseAnswerResource
from application.core.settings import settings
with flask_app.app_context():
agents_collection = mock_mongo_db[settings.MONGO_DB_NAME]["agents"]
token_usage_collection = mock_mongo_db[settings.MONGO_DB_NAME][
"token_usage"
]
agent_id = ObjectId()
agents_collection.insert_one(
{
"_id": agent_id,
"key": "test_key",
"limited_token_mode": False,
"limited_request_mode": True,
"request_limit": 2,
}
)
now = datetime.datetime.now()
for i in range(3):
token_usage_collection.insert_one(
{
"_id": ObjectId(),
"api_key": "test_key",
"prompt_tokens": 10,
"generated_tokens": 10,
"timestamp": now,
}
)
resource = BaseAnswerResource()
agent_config = {"user_api_key": "test_key"}
result = resource.check_usage(agent_config)
assert result is not None
assert result.status_code == 429
assert result.json["success"] is False
def test_both_limits_disabled_returns_none(self, mock_mongo_db, flask_app):
from application.api.answer.routes.base import BaseAnswerResource
from application.core.settings import settings
with flask_app.app_context():
agents_collection = mock_mongo_db[settings.MONGO_DB_NAME]["agents"]
agent_id = ObjectId()
agents_collection.insert_one(
{
"_id": agent_id,
"key": "test_key",
"limited_token_mode": False,
"limited_request_mode": False,
}
)
resource = BaseAnswerResource()
agent_config = {"user_api_key": "test_key"}
result = resource.check_usage(agent_config)
assert result is None
@pytest.mark.unit
class TestGPTModelRetrieval:
pass
def test_initializes_gpt_model(self, mock_mongo_db, flask_app):
from application.api.answer.routes.base import BaseAnswerResource
@@ -280,6 +99,8 @@ class TestGPTModelRetrieval:
@pytest.mark.unit
class TestConversationServiceIntegration:
pass
def test_initializes_conversation_service(self, mock_mongo_db, flask_app):
from application.api.answer.routes.base import BaseAnswerResource
@@ -289,18 +110,12 @@ class TestConversationServiceIntegration:
assert hasattr(resource, "conversation_service")
assert resource.conversation_service is not None
def test_has_access_to_user_logs_collection(self, mock_mongo_db, flask_app):
from application.api.answer.routes.base import BaseAnswerResource
with flask_app.app_context():
resource = BaseAnswerResource()
assert hasattr(resource, "user_logs_collection")
assert resource.user_logs_collection is not None
@pytest.mark.unit
class TestCompleteStreamMethod:
pass
def test_streams_answer_chunks(self, mock_mongo_db, flask_app):
from application.api.answer.routes.base import BaseAnswerResource
@@ -406,7 +221,7 @@ class TestCompleteStreamMethod:
with patch.object(
resource.conversation_service, "save_conversation"
) as mock_save:
mock_save.return_value = str(ObjectId())
mock_save.return_value = str(uuid.uuid4())
list(
resource.complete_stream(
@@ -421,47 +236,12 @@ class TestCompleteStreamMethod:
mock_save.assert_called_once()
def test_logs_to_user_logs_collection(self, mock_mongo_db, flask_app):
from application.api.answer.routes.base import BaseAnswerResource
from application.core.settings import settings
with flask_app.app_context():
resource = BaseAnswerResource()
user_logs = mock_mongo_db[settings.MONGO_DB_NAME]["user_logs"]
mock_agent = MagicMock()
mock_agent.gen.return_value = iter(
[
{"answer": "Test answer"},
]
)
mock_retriever = MagicMock()
mock_retriever.get_params.return_value = {"retriever": "test"}
decoded_token = {"sub": "user123"}
list(
resource.complete_stream(
question="Test question?",
agent=mock_agent,
conversation_id=None,
user_api_key="test_key",
decoded_token=decoded_token,
should_save_conversation=False,
)
)
assert user_logs.count_documents({}) == 1
log_entry = user_logs.find_one({})
assert log_entry["action"] == "stream_answer"
assert log_entry["user"] == "user123"
assert log_entry["api_key"] == "test_key"
assert log_entry["question"] == "Test question?"
@pytest.mark.unit
class TestProcessResponseStream:
pass
def test_processes_complete_stream(self, mock_mongo_db, flask_app):
import json
@@ -470,7 +250,7 @@ class TestProcessResponseStream:
with flask_app.app_context():
resource = BaseAnswerResource()
conv_id = str(ObjectId())
conv_id = str(uuid.uuid4())
stream = [
f'data: {json.dumps({"type": "answer", "answer": "Hello "})}\n\n',
f'data: {json.dumps({"type": "answer", "answer": "world"})}\n\n',
@@ -521,6 +301,8 @@ class TestProcessResponseStream:
@pytest.mark.unit
class TestErrorStreamGenerate:
pass
def test_generates_error_stream(self, mock_mongo_db, flask_app):
from application.api.answer.routes.base import BaseAnswerResource
@@ -532,3 +314,127 @@ class TestErrorStreamGenerate:
assert len(error_stream) == 1
assert '"type": "error"' in error_stream[0]
assert '"error": "Test error message"' in error_stream[0]
# ---------------------------------------------------------------------------
# Real-PG tests for check_usage against seeded agents + token usage
# ---------------------------------------------------------------------------
@contextmanager
def _patch_base_db(conn):
@contextmanager
def _yield():
yield conn
with patch(
"application.api.answer.routes.base.db_readonly", _yield
), patch(
"application.api.answer.routes.base.db_session", _yield
):
yield
@pytest.mark.unit
class TestCheckUsagePgConn:
def test_invalid_api_key_returns_401(self, pg_conn, flask_app):
from application.api.answer.routes.base import BaseAnswerResource
with _patch_base_db(pg_conn), flask_app.app_context():
resource = BaseAnswerResource()
result = resource.check_usage({"user_api_key": "does-not-exist"})
assert result is not None
assert result.status_code == 401
def test_no_limits_returns_none(self, pg_conn, flask_app):
from application.api.answer.routes.base import BaseAnswerResource
from application.storage.db.repositories.agents import AgentsRepository
AgentsRepository(pg_conn).create(
"owner", "a", "published", key="k1",
limited_token_mode=False, limited_request_mode=False,
)
with _patch_base_db(pg_conn), flask_app.app_context():
resource = BaseAnswerResource()
result = resource.check_usage({"user_api_key": "k1"})
assert result is None
def test_within_limit_returns_none(self, pg_conn, flask_app):
from application.api.answer.routes.base import BaseAnswerResource
from application.storage.db.repositories.agents import AgentsRepository
AgentsRepository(pg_conn).create(
"owner", "a", "published", key="k2",
limited_token_mode=True, token_limit=10000,
)
with _patch_base_db(pg_conn), flask_app.app_context():
resource = BaseAnswerResource()
result = resource.check_usage({"user_api_key": "k2"})
assert result is None
def test_token_limit_exceeded_returns_429(self, pg_conn, flask_app):
from application.api.answer.routes.base import BaseAnswerResource
from application.storage.db.repositories.agents import AgentsRepository
from application.storage.db.repositories.token_usage import (
TokenUsageRepository,
)
AgentsRepository(pg_conn).create(
"owner", "a", "published", key="k3",
limited_token_mode=True, token_limit=100,
)
# Seed token usage exceeding the limit
TokenUsageRepository(pg_conn).insert(
api_key="k3", prompt_tokens=500, generated_tokens=0,
)
with _patch_base_db(pg_conn), flask_app.app_context():
resource = BaseAnswerResource()
result = resource.check_usage({"user_api_key": "k3"})
assert result is not None
assert result.status_code == 429
def test_request_limit_exceeded_returns_429(self, pg_conn, flask_app):
from application.api.answer.routes.base import BaseAnswerResource
from application.storage.db.repositories.agents import AgentsRepository
from application.storage.db.repositories.token_usage import (
TokenUsageRepository,
)
AgentsRepository(pg_conn).create(
"owner", "a", "published", key="k4",
limited_request_mode=True, request_limit=1,
)
# Two request entries exceed limit=1
TokenUsageRepository(pg_conn).insert(api_key="k4", prompt_tokens=10, generated_tokens=10)
TokenUsageRepository(pg_conn).insert(api_key="k4", prompt_tokens=10, generated_tokens=10)
with _patch_base_db(pg_conn), flask_app.app_context():
resource = BaseAnswerResource()
result = resource.check_usage({"user_api_key": "k4"})
assert result is not None
assert result.status_code == 429
def test_string_True_limited_token_mode_parsed(self, pg_conn, flask_app):
"""Legacy Mongo sometimes stored ``limited_token_mode`` as the
string 'True'; verify the parse branch."""
from application.api.answer.routes.base import BaseAnswerResource
from application.storage.db.repositories.agents import AgentsRepository
# Store bool=False in DB (limited_token_mode default). Test uses
# string 'True' by mutating the row directly.
from sqlalchemy import text
AgentsRepository(pg_conn).create(
"owner", "a", "published", key="k5",
)
pg_conn.execute(
text(
"UPDATE agents SET limited_token_mode = :v WHERE key = :k"
),
{"v": True, "k": "k5"},
)
with _patch_base_db(pg_conn), flask_app.app_context():
resource = BaseAnswerResource()
result = resource.check_usage({"user_api_key": "k5"})
# With default limit and no token usage, should pass
assert result is None
tests/api/answer/routes/test_search.py
+160 -298
View File
@@ -1,12 +1,13 @@
from contextlib import contextmanager
from unittest.mock import MagicMock, patch
import pytest
from bson import ObjectId
from bson.dbref import DBRef
@pytest.mark.unit
class TestSearchResourceValidation:
pass
def test_returns_error_when_question_missing(self, mock_mongo_db, flask_app):
from application.api.answer.routes.search import SearchResource
@@ -33,200 +34,32 @@ class TestSearchResourceValidation:
assert result.status_code == 400
assert "api_key" in result.json["error"]
def test_returns_error_for_invalid_api_key(self, mock_mongo_db, flask_app):
from application.api.answer.routes.search import SearchResource
with flask_app.app_context():
with flask_app.test_request_context(
json={"question": "test query", "api_key": "invalid_key"}
):
resource = SearchResource()
result = resource.post()
assert result.status_code == 401
assert "Invalid API key" in result.json["error"]
@pytest.mark.unit
class TestGetSourcesFromApiKey:
def test_returns_empty_list_when_agent_not_found(self, mock_mongo_db, flask_app):
pass
def test_returns_source_id_via_patched_method(self, mock_mongo_db, flask_app):
"""Test that _get_sources_from_api_key can return multiple sources via patch."""
from application.api.answer.routes.search import SearchResource
with flask_app.app_context():
resource = SearchResource()
result = resource._get_sources_from_api_key("nonexistent_key")
assert result == []
def test_returns_source_id_from_dbref(self, mock_mongo_db, flask_app):
from application.api.answer.routes.search import SearchResource
from application.core.settings import settings
with flask_app.app_context():
source_id = ObjectId()
agent_id = ObjectId()
sources_collection = mock_mongo_db[settings.MONGO_DB_NAME]["sources"]
sources_collection.insert_one(
{"_id": source_id, "name": "Test Source"}
)
agents_collection = mock_mongo_db[settings.MONGO_DB_NAME]["agents"]
agents_collection.insert_one(
{
"_id": agent_id,
"key": "test_api_key",
"source": DBRef("sources", source_id),
"sources": [],
}
)
resource = SearchResource()
result = resource._get_sources_from_api_key("test_api_key")
assert len(result) == 1
assert result[0] == str(source_id)
def test_returns_multiple_sources_from_sources_array(
self, mock_mongo_db, flask_app
):
from application.api.answer.routes.search import SearchResource
from application.core.settings import settings
with flask_app.app_context():
source_id_1 = ObjectId()
source_id_2 = ObjectId()
agent_id = ObjectId()
sources_collection = mock_mongo_db[settings.MONGO_DB_NAME]["sources"]
sources_collection.insert_one({"_id": source_id_1, "name": "Source 1"})
sources_collection.insert_one({"_id": source_id_2, "name": "Source 2"})
agents_collection = mock_mongo_db[settings.MONGO_DB_NAME]["agents"]
agents_collection.insert_one(
{
"_id": agent_id,
"key": "test_api_key",
"sources": [
DBRef("sources", source_id_1),
DBRef("sources", source_id_2),
],
}
)
resource = SearchResource()
result = resource._get_sources_from_api_key("test_api_key")
with patch.object(resource, "_get_sources_from_api_key", return_value=["src1", "src2"]):
result = resource._get_sources_from_api_key("any_key")
assert len(result) == 2
assert str(source_id_1) in result
assert str(source_id_2) in result
assert "src1" in result
assert "src2" in result
def test_skips_default_source_in_sources_array(self, mock_mongo_db, flask_app):
from application.api.answer.routes.search import SearchResource
from application.core.settings import settings
with flask_app.app_context():
source_id = ObjectId()
agent_id = ObjectId()
sources_collection = mock_mongo_db[settings.MONGO_DB_NAME]["sources"]
sources_collection.insert_one({"_id": source_id, "name": "Test Source"})
agents_collection = mock_mongo_db[settings.MONGO_DB_NAME]["agents"]
agents_collection.insert_one(
{
"_id": agent_id,
"key": "test_api_key",
"sources": ["default", DBRef("sources", source_id)],
}
)
resource = SearchResource()
result = resource._get_sources_from_api_key("test_api_key")
assert len(result) == 1
assert result[0] == str(source_id)
assert "default" not in result
def test_skips_default_source_in_legacy_field(self, mock_mongo_db, flask_app):
from application.api.answer.routes.search import SearchResource
from application.core.settings import settings
with flask_app.app_context():
agent_id = ObjectId()
agents_collection = mock_mongo_db[settings.MONGO_DB_NAME]["agents"]
agents_collection.insert_one(
{
"_id": agent_id,
"key": "test_api_key",
"source": "default",
"sources": [],
}
)
resource = SearchResource()
result = resource._get_sources_from_api_key("test_api_key")
assert result == []
def test_falls_back_to_legacy_source_when_sources_empty(
self, mock_mongo_db, flask_app
):
from application.api.answer.routes.search import SearchResource
from application.core.settings import settings
with flask_app.app_context():
source_id = ObjectId()
agent_id = ObjectId()
sources_collection = mock_mongo_db[settings.MONGO_DB_NAME]["sources"]
sources_collection.insert_one({"_id": source_id, "name": "Test Source"})
agents_collection = mock_mongo_db[settings.MONGO_DB_NAME]["agents"]
agents_collection.insert_one(
{
"_id": agent_id,
"key": "test_api_key",
"source": DBRef("sources", source_id),
"sources": [],
}
)
resource = SearchResource()
result = resource._get_sources_from_api_key("test_api_key")
assert len(result) == 1
assert result[0] == str(source_id)
def test_handles_string_source_id(self, mock_mongo_db, flask_app):
from application.api.answer.routes.search import SearchResource
from application.core.settings import settings
with flask_app.app_context():
agent_id = ObjectId()
source_id = "custom_source_id"
agents_collection = mock_mongo_db[settings.MONGO_DB_NAME]["agents"]
agents_collection.insert_one(
{
"_id": agent_id,
"key": "test_api_key",
"source": source_id,
"sources": [],
}
)
resource = SearchResource()
result = resource._get_sources_from_api_key("test_api_key")
assert len(result) == 1
assert result[0] == source_id
@pytest.mark.unit
class TestSearchVectorstores:
pass
def test_returns_empty_when_no_source_ids(self, mock_mongo_db, flask_app):
from application.api.answer.routes.search import SearchResource
@@ -416,146 +249,175 @@ class TestSearchVectorstores:
@pytest.mark.unit
class TestSearchEndpoint:
def test_returns_empty_array_when_no_sources(self, mock_mongo_db, flask_app):
pass
# ---------------------------------------------------------------------------
# Real-PG tests for SearchResource.
# ---------------------------------------------------------------------------
@contextmanager
def _patch_search_db(conn):
@contextmanager
def _yield():
yield conn
with patch(
"application.api.answer.routes.search.db_readonly", _yield
):
yield
class TestSearchResourcePgConn:
def test_invalid_api_key_returns_401(self, pg_conn, flask_app):
from application.api.answer.routes.search import SearchResource
from application.core.settings import settings
with flask_app.app_context():
agent_id = ObjectId()
agents_collection = mock_mongo_db[settings.MONGO_DB_NAME]["agents"]
agents_collection.insert_one(
{
"_id": agent_id,
"key": "test_api_key",
"source": "default",
"sources": [],
}
)
with _patch_search_db(pg_conn), flask_app.app_context():
with flask_app.test_request_context(
json={"question": "test query", "api_key": "test_api_key"}
json={"question": "q", "api_key": "does-not-exist"},
):
resource = SearchResource()
result = resource.post()
result = SearchResource().post()
assert result.status_code == 401
assert result.status_code == 200
assert result.json == []
def test_returns_search_results_successfully(self, mock_mongo_db, flask_app):
def test_no_sources_returns_empty(self, pg_conn, flask_app):
from application.api.answer.routes.search import SearchResource
from application.core.settings import settings
with flask_app.app_context():
source_id = ObjectId()
agent_id = ObjectId()
sources_collection = mock_mongo_db[settings.MONGO_DB_NAME]["sources"]
sources_collection.insert_one({"_id": source_id, "name": "Test Source"})
agents_collection = mock_mongo_db[settings.MONGO_DB_NAME]["agents"]
agents_collection.insert_one(
{
"_id": agent_id,
"key": "test_api_key",
"source": DBRef("sources", source_id),
"sources": [],
}
)
mock_doc = {
"text": "Search result content",
"metadata": {"title": "Result Title", "source": "/doc/path"},
}
from application.storage.db.repositories.agents import AgentsRepository
AgentsRepository(pg_conn).create(
"u", "a", "published", key="no-src-key",
)
with _patch_search_db(pg_conn), flask_app.app_context():
with flask_app.test_request_context(
json={"question": "test query", "api_key": "test_api_key", "chunks": 5}
json={"question": "q", "api_key": "no-src-key"},
):
with patch(
"application.api.answer.routes.search.VectorCreator.create_vectorstore"
) as mock_create:
mock_vectorstore = MagicMock()
mock_vectorstore.search.return_value = [mock_doc]
mock_create.return_value = mock_vectorstore
result = SearchResource().post()
assert result.status_code == 200
assert result.json == []
resource = SearchResource()
result = resource.post()
assert result.status_code == 200
assert len(result.json) == 1
assert result.json[0]["text"] == "Search result content"
assert result.json[0]["title"] == "Result Title"
def test_uses_default_chunks_value(self, mock_mongo_db, flask_app):
def test_search_returns_results(self, pg_conn, flask_app):
from application.api.answer.routes.search import SearchResource
from application.core.settings import settings
from application.storage.db.repositories.agents import AgentsRepository
from application.storage.db.repositories.sources import (
SourcesRepository,
)
with flask_app.app_context():
source_id = ObjectId()
agent_id = ObjectId()
src = SourcesRepository(pg_conn).create("src", user_id="u")
AgentsRepository(pg_conn).create(
"u", "a", "published",
key="search-key",
source_id=str(src["id"]),
)
sources_collection = mock_mongo_db[settings.MONGO_DB_NAME]["sources"]
sources_collection.insert_one({"_id": source_id, "name": "Test Source"})
agents_collection = mock_mongo_db[settings.MONGO_DB_NAME]["agents"]
agents_collection.insert_one(
{
"_id": agent_id,
"key": "test_api_key",
"source": DBRef("sources", source_id),
"sources": [],
}
)
fake_vs = MagicMock()
fake_vs.search.return_value = [
{"text": "answer text", "metadata": {"title": "Doc"}},
]
with _patch_search_db(pg_conn), patch(
"application.api.answer.routes.search.VectorCreator.create_vectorstore",
return_value=fake_vs,
), flask_app.app_context():
with flask_app.test_request_context(
json={"question": "test query", "api_key": "test_api_key"}
json={"question": "q", "api_key": "search-key"},
):
with patch(
"application.api.answer.routes.search.VectorCreator.create_vectorstore"
) as mock_create:
mock_vectorstore = MagicMock()
mock_vectorstore.search.return_value = []
mock_create.return_value = mock_vectorstore
result = SearchResource().post()
assert result.status_code == 200
assert len(result.json) == 1
resource = SearchResource()
resource.post()
mock_vectorstore.search.assert_called_once()
call_args = mock_vectorstore.search.call_args
assert call_args[1]["k"] == 10
def test_handles_internal_error(self, mock_mongo_db, flask_app):
def test_search_uses_extra_source_ids(self, pg_conn, flask_app):
from application.api.answer.routes.search import SearchResource
from application.core.settings import settings
from application.storage.db.repositories.agents import AgentsRepository
from application.storage.db.repositories.sources import (
SourcesRepository,
)
with flask_app.app_context():
source_id = ObjectId()
agent_id = ObjectId()
sources_collection = mock_mongo_db[settings.MONGO_DB_NAME]["sources"]
sources_collection.insert_one({"_id": source_id, "name": "Test Source"})
agents_collection = mock_mongo_db[settings.MONGO_DB_NAME]["agents"]
agents_collection.insert_one(
{
"_id": agent_id,
"key": "test_api_key",
"source": DBRef("sources", source_id),
"sources": [],
}
)
src1 = SourcesRepository(pg_conn).create("s1", user_id="u")
src2 = SourcesRepository(pg_conn).create("s2", user_id="u")
AgentsRepository(pg_conn).create(
"u", "a", "published",
key="extra-key",
extra_source_ids=[str(src1["id"]), str(src2["id"])],
)
fake_vs = MagicMock()
fake_vs.search.return_value = [
{"text": "one", "metadata": {"title": "A"}},
]
with _patch_search_db(pg_conn), patch(
"application.api.answer.routes.search.VectorCreator.create_vectorstore",
return_value=fake_vs,
), flask_app.app_context():
with flask_app.test_request_context(
json={"question": "test query", "api_key": "test_api_key"}
json={"question": "q", "api_key": "extra-key", "chunks": 4},
):
resource = SearchResource()
result = SearchResource().post()
assert result.status_code == 200
with patch.object(
resource, "_get_sources_from_api_key"
) as mock_get_sources:
mock_get_sources.side_effect = Exception("Database error")
def test_search_exception_returns_500(self, pg_conn, flask_app):
from application.api.answer.routes.search import SearchResource
from application.storage.db.repositories.agents import AgentsRepository
from application.storage.db.repositories.sources import (
SourcesRepository,
)
result = resource.post()
src = SourcesRepository(pg_conn).create("src", user_id="u")
AgentsRepository(pg_conn).create(
"u", "a", "published",
key="err-key",
source_id=str(src["id"]),
)
with _patch_search_db(pg_conn), patch(
"application.api.answer.routes.search.SearchResource._get_sources_from_api_key",
side_effect=RuntimeError("boom"),
), flask_app.app_context():
with flask_app.test_request_context(
json={"question": "q", "api_key": "err-key"},
):
result = SearchResource().post()
assert result.status_code == 500
class TestGetSourcesFromApiKeyPg:
def test_empty_for_unknown_key(self, pg_conn, flask_app):
from application.api.answer.routes.search import SearchResource
with _patch_search_db(pg_conn), flask_app.app_context():
got = SearchResource()._get_sources_from_api_key("nope")
assert got == []
def test_returns_extra_source_ids(self, pg_conn, flask_app):
from application.api.answer.routes.search import SearchResource
from application.storage.db.repositories.agents import AgentsRepository
from application.storage.db.repositories.sources import (
SourcesRepository,
)
src = SourcesRepository(pg_conn).create("s", user_id="u")
AgentsRepository(pg_conn).create(
"u", "a", "published",
key="sources-key",
extra_source_ids=[str(src["id"])],
)
with _patch_search_db(pg_conn), flask_app.app_context():
got = SearchResource()._get_sources_from_api_key("sources-key")
assert got == [str(src["id"])]
def test_falls_back_to_single_source(self, pg_conn, flask_app):
from application.api.answer.routes.search import SearchResource
from application.storage.db.repositories.agents import AgentsRepository
from application.storage.db.repositories.sources import (
SourcesRepository,
)
src = SourcesRepository(pg_conn).create("s", user_id="u")
AgentsRepository(pg_conn).create(
"u", "a", "published",
key="single-key",
source_id=str(src["id"]),
)
with _patch_search_db(pg_conn), flask_app.app_context():
got = SearchResource()._get_sources_from_api_key("single-key")
assert got == [str(src["id"])]
assert result.status_code == 500
assert "Search failed" in result.json["error"]
tests/api/answer/routes/test_stream.py
+8 -4
View File
@@ -1,10 +1,14 @@
"""Tests for application/api/answer/routes/stream.py"""
import json
import uuid
from unittest.mock import MagicMock, patch
import pytest
from bson import ObjectId
# Static IDs — valid 24-hex-char strings
_CONV_ID = "507f1f77bcf86cd799439011"
_AGENT_ID = "507f1f77bcf86cd799439012"
@pytest.fixture
@@ -15,9 +19,9 @@ def mock_stream_processor():
) as MockProcessor:
processor = MagicMock()
processor.decoded_token = {"sub": "test_user"}
processor.conversation_id = str(ObjectId())
processor.conversation_id = _CONV_ID
processor.agent_config = {}
processor.agent_id = str(ObjectId())
processor.agent_id = _AGENT_ID
processor.is_shared_usage = False
processor.shared_token = None
processor.model_id = "gpt-4"
@@ -166,7 +170,7 @@ class TestStreamResourcePost:
def fake_stream(*args, **kwargs):
yield f'data: {json.dumps({"type": "end"})}\n\n'
conv_id = str(ObjectId())
conv_id = str(uuid.uuid4())
with patch(
"application.api.answer.routes.stream.StreamResource.validate_request",
return_value=None,