wehub-resource-sync
fed8b2eed7
Build and push multi-arch DocsGPT Docker image / build (linux/amd64, ubuntu-latest, amd64) (push) Has been cancelled
Backend release / release (push) Has been cancelled
Bandit Security Scan / bandit_scan (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / manifest (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / build (linux/amd64, ubuntu-latest, amd64) (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / manifest (push) Has been cancelled
Python linting / ruff (push) Has been cancelled
Run python tests with pytest / Run tests and count coverage (3.12) (push) Has been cancelled
React Widget Build / build (push) Has been cancelled
701 lines
25 KiB
Python
701 lines
25 KiB
Python
import uuid
|
|
from contextlib import contextmanager
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestBaseAnswerValidation:
|
|
pass
|
|
|
|
def test_validate_request_passes_with_required_fields(
|
|
self, mock_mongo_db, flask_app
|
|
):
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
data = {"question": "What is Python?"}
|
|
|
|
result = resource.validate_request(data)
|
|
|
|
assert result is None
|
|
|
|
def test_validate_request_fails_without_question(self, mock_mongo_db, flask_app):
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
data = {}
|
|
|
|
result = resource.validate_request(data)
|
|
|
|
assert result is not None
|
|
assert result.status_code == 400
|
|
assert "question" in result.json["message"].lower()
|
|
|
|
def test_validate_with_conversation_id_required(self, mock_mongo_db, flask_app):
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
data = {"question": "Test"}
|
|
|
|
result = resource.validate_request(data, require_conversation_id=True)
|
|
|
|
assert result is not None
|
|
assert result.status_code == 400
|
|
assert "conversation_id" in result.json["message"].lower()
|
|
|
|
def test_validate_passes_with_all_required_fields(self, mock_mongo_db, flask_app):
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
data = {"question": "Test", "conversation_id": str(uuid.uuid4())}
|
|
|
|
result = resource.validate_request(data, require_conversation_id=True)
|
|
|
|
assert result is None
|
|
|
|
|
|
@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
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
agent_config = {}
|
|
|
|
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
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
|
|
assert hasattr(resource, "default_model_id")
|
|
assert resource.default_model_id is not None
|
|
|
|
|
|
@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
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
|
|
assert hasattr(resource, "conversation_service")
|
|
assert resource.conversation_service 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
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
|
|
mock_agent = MagicMock()
|
|
mock_agent.gen.return_value = iter(
|
|
[
|
|
{"answer": "Hello "},
|
|
{"answer": "world!"},
|
|
]
|
|
)
|
|
|
|
decoded_token = {"sub": "user123"}
|
|
|
|
stream = list(
|
|
resource.complete_stream(
|
|
question="Test question",
|
|
agent=mock_agent,
|
|
conversation_id=None,
|
|
user_api_key=None,
|
|
decoded_token=decoded_token,
|
|
should_persist=False,
|
|
)
|
|
)
|
|
|
|
answer_chunks = [s for s in stream if '"type": "answer"' in s]
|
|
assert len(answer_chunks) == 2
|
|
assert '"answer": "Hello "' in answer_chunks[0]
|
|
assert '"answer": "world!"' in answer_chunks[1]
|
|
|
|
def test_streams_sources(self, mock_mongo_db, flask_app):
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
|
|
mock_agent = MagicMock()
|
|
mock_agent.gen.return_value = iter(
|
|
[
|
|
{"answer": "Test answer"},
|
|
{"sources": [{"title": "doc1.txt", "text": "x" * 200}]},
|
|
]
|
|
)
|
|
|
|
decoded_token = {"sub": "user123"}
|
|
|
|
stream = list(
|
|
resource.complete_stream(
|
|
question="Test?",
|
|
agent=mock_agent,
|
|
conversation_id=None,
|
|
user_api_key=None,
|
|
decoded_token=decoded_token,
|
|
should_persist=False,
|
|
)
|
|
)
|
|
|
|
source_chunks = [s for s in stream if '"type": "source"' in s]
|
|
assert len(source_chunks) == 1
|
|
assert '"title": "doc1.txt"' in source_chunks[0]
|
|
|
|
def test_handles_error_during_streaming(self, mock_mongo_db, flask_app):
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
|
|
mock_agent = MagicMock()
|
|
mock_agent.gen.side_effect = Exception("Test error")
|
|
|
|
decoded_token = {"sub": "user123"}
|
|
|
|
stream = list(
|
|
resource.complete_stream(
|
|
question="Test?",
|
|
agent=mock_agent,
|
|
conversation_id=None,
|
|
user_api_key=None,
|
|
decoded_token=decoded_token,
|
|
should_persist=False,
|
|
)
|
|
)
|
|
|
|
assert any('"type": "error"' in s for s in stream)
|
|
|
|
def test_user_facing_error_is_not_sanitized(self, mock_mongo_db, flask_app):
|
|
"""A user_facing error (e.g. an artifact-quota notice) streams verbatim.
|
|
|
|
Without the flag, sanitize_api_error substring-matches "quota" and rewrites the
|
|
message into a misleading rate-limit notice.
|
|
"""
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
|
|
mock_agent = MagicMock()
|
|
mock_agent.gen.return_value = iter(
|
|
[
|
|
{
|
|
"type": "error",
|
|
"user_facing": True,
|
|
"error": "This run's input documents exceed your artifact storage quota.",
|
|
}
|
|
]
|
|
)
|
|
|
|
stream = list(
|
|
resource.complete_stream(
|
|
question="Test?",
|
|
agent=mock_agent,
|
|
conversation_id=None,
|
|
user_api_key=None,
|
|
decoded_token={"sub": "user123"},
|
|
should_persist=False,
|
|
)
|
|
)
|
|
|
|
error_chunks = [s for s in stream if '"type": "error"' in s]
|
|
assert error_chunks
|
|
assert "artifact storage quota" in error_chunks[0]
|
|
assert "Rate limit exceeded" not in error_chunks[0]
|
|
|
|
def test_notice_is_forwarded_verbatim_and_not_an_error(self, mock_mongo_db, flask_app):
|
|
"""A non-fatal ``notice`` streams through as a notice, never as an error.
|
|
|
|
A ``notice`` (e.g. some workflow input documents were dropped) must not be
|
|
emitted as ``type: error`` -- the client treats an error event as terminal and
|
|
disables reconnect -- and its text must not be run through sanitize_api_error.
|
|
"""
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
|
|
mock_agent = MagicMock()
|
|
mock_agent.gen.return_value = iter(
|
|
[{"type": "notice", "notice": "big.txt exceeds the per-file size limit"}]
|
|
)
|
|
|
|
stream = list(
|
|
resource.complete_stream(
|
|
question="Test?",
|
|
agent=mock_agent,
|
|
conversation_id=None,
|
|
user_api_key=None,
|
|
decoded_token={"sub": "user123"},
|
|
should_persist=False,
|
|
)
|
|
)
|
|
|
|
notice_chunks = [s for s in stream if '"type": "notice"' in s]
|
|
assert notice_chunks
|
|
assert "big.txt exceeds the per-file size limit" in notice_chunks[0]
|
|
# Crucially, it is not surfaced as an error event.
|
|
assert not [s for s in stream if '"type": "error"' in s]
|
|
|
|
def test_non_user_facing_error_is_sanitized(self, mock_mongo_db, flask_app):
|
|
"""A raw error without the flag is still routed through sanitize_api_error."""
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
|
|
mock_agent = MagicMock()
|
|
mock_agent.gen.return_value = iter(
|
|
[{"type": "error", "error": "OpenAI 429: quota exceeded for this key"}]
|
|
)
|
|
|
|
stream = list(
|
|
resource.complete_stream(
|
|
question="Test?",
|
|
agent=mock_agent,
|
|
conversation_id=None,
|
|
user_api_key=None,
|
|
decoded_token={"sub": "user123"},
|
|
should_persist=False,
|
|
)
|
|
)
|
|
|
|
error_chunks = [s for s in stream if '"type": "error"' in s]
|
|
assert error_chunks
|
|
assert "Rate limit exceeded" in error_chunks[0]
|
|
|
|
def test_saves_conversation_when_enabled(self, mock_mongo_db, flask_app):
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
|
|
mock_agent = MagicMock()
|
|
mock_agent.gen.return_value = iter(
|
|
[
|
|
{"answer": "Test answer"},
|
|
]
|
|
)
|
|
|
|
decoded_token = {"sub": "user123"}
|
|
|
|
# The fresh-question path now reserves a row before agent.gen()
|
|
# and calls finalize_message at end of stream — assert both fire.
|
|
with patch.object(
|
|
resource.conversation_service, "save_user_question"
|
|
) as mock_reserve, patch.object(
|
|
resource.conversation_service, "finalize_message"
|
|
) as mock_finalize:
|
|
mock_reserve.return_value = {
|
|
"conversation_id": str(uuid.uuid4()),
|
|
"message_id": str(uuid.uuid4()),
|
|
"request_id": "req-1",
|
|
}
|
|
mock_finalize.return_value = True
|
|
|
|
list(
|
|
resource.complete_stream(
|
|
question="Test?",
|
|
agent=mock_agent,
|
|
conversation_id=None,
|
|
user_api_key=None,
|
|
decoded_token=decoded_token,
|
|
should_persist=True,
|
|
)
|
|
)
|
|
|
|
mock_reserve.assert_called_once()
|
|
mock_finalize.assert_called_once()
|
|
|
|
def test_tool_executor_conversation_id_set_after_reserve(
|
|
self, mock_mongo_db, flask_app,
|
|
):
|
|
"""Regression: ``save_user_question`` may mint a fresh
|
|
``conversation_id`` (first turn). The propagation MUST land on
|
|
``agent.tool_executor.conversation_id`` BEFORE ``agent.gen`` runs,
|
|
so tools needing a conversation home (``scheduler`` in an agentless
|
|
chat) see it on the very first call.
|
|
"""
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
|
|
fresh_conv_id = str(uuid.uuid4())
|
|
seen_conv_id_on_gen: dict = {}
|
|
|
|
mock_agent = MagicMock()
|
|
tool_executor = MagicMock()
|
|
# Start with no conversation_id — the propagation must set it.
|
|
tool_executor.conversation_id = None
|
|
mock_agent.tool_executor = tool_executor
|
|
|
|
def _gen(**_kwargs):
|
|
# Capture the executor's id at the exact moment gen runs;
|
|
# this is what tools see when called from the agent loop.
|
|
seen_conv_id_on_gen["value"] = (
|
|
mock_agent.tool_executor.conversation_id
|
|
)
|
|
yield {"answer": "ok"}
|
|
|
|
mock_agent.gen.side_effect = _gen
|
|
mock_agent.gen.return_value = None # use side_effect instead
|
|
|
|
with patch.object(
|
|
resource.conversation_service, "save_user_question"
|
|
) as mock_reserve, patch.object(
|
|
resource.conversation_service, "finalize_message",
|
|
return_value=True,
|
|
):
|
|
mock_reserve.return_value = {
|
|
"conversation_id": fresh_conv_id,
|
|
"message_id": str(uuid.uuid4()),
|
|
"request_id": "req-prop",
|
|
}
|
|
|
|
list(
|
|
resource.complete_stream(
|
|
question="schedule something",
|
|
agent=mock_agent,
|
|
conversation_id=None, # caller had no conv yet
|
|
user_api_key=None,
|
|
decoded_token={"sub": "user-prop"},
|
|
should_persist=True,
|
|
)
|
|
)
|
|
|
|
# The fresh id reserved by save_user_question must reach the
|
|
# tool_executor before agent.gen consumes it.
|
|
assert seen_conv_id_on_gen["value"] == fresh_conv_id
|
|
assert tool_executor.conversation_id == fresh_conv_id
|
|
|
|
def _run_paused(self, resource, pending_calls):
|
|
"""Drive complete_stream into its paused branch with the given pending
|
|
tool calls, mocking out the WAL row and continuation save."""
|
|
agent = MagicMock()
|
|
agent.gen.return_value = iter(
|
|
[{"type": "tool_calls_pending", "data": {"pending_tool_calls": pending_calls}}]
|
|
)
|
|
agent._pending_continuation = {
|
|
"messages": [],
|
|
"pending_tool_calls": pending_calls,
|
|
"tools_dict": {},
|
|
}
|
|
# Make the WAL reservation no-op so we stay off the journal/DB.
|
|
resource.conversation_service = MagicMock()
|
|
resource.conversation_service.save_user_question.side_effect = Exception("skip")
|
|
list(
|
|
resource.complete_stream(
|
|
question="Do the test",
|
|
agent=agent,
|
|
conversation_id="conv-1",
|
|
user_api_key=None,
|
|
decoded_token={"sub": "user123"},
|
|
should_persist=True,
|
|
)
|
|
)
|
|
|
|
def test_paused_skips_notification_for_client_execution(
|
|
self, mock_mongo_db, flask_app
|
|
):
|
|
"""A pure ``requires_client_execution`` pause must NOT publish a
|
|
``tool.approval.required`` event — the client resolves it, so the
|
|
notification would be non-actionable noise."""
|
|
from application.api.answer.routes import base as base_mod
|
|
|
|
with flask_app.app_context(), patch.object(
|
|
base_mod, "publish_user_event"
|
|
) as published, patch.object(
|
|
base_mod, "ContinuationService", MagicMock
|
|
):
|
|
self._run_paused(
|
|
base_mod.BaseAnswerResource(),
|
|
[
|
|
{
|
|
"call_id": "c1",
|
|
"name": "create_file",
|
|
"tool_name": "create_file",
|
|
"action_name": "create_file",
|
|
"pause_type": "requires_client_execution",
|
|
}
|
|
],
|
|
)
|
|
published.assert_not_called()
|
|
|
|
def test_paused_publishes_notification_only_for_awaiting_approval(
|
|
self, mock_mongo_db, flask_app
|
|
):
|
|
"""A pause with an ``awaiting_approval`` call publishes once, and the
|
|
payload surfaces only the approval call (not the client-side one)."""
|
|
from application.api.answer.routes import base as base_mod
|
|
|
|
with flask_app.app_context(), patch.object(
|
|
base_mod, "publish_user_event"
|
|
) as published, patch.object(
|
|
base_mod, "ContinuationService", MagicMock
|
|
):
|
|
self._run_paused(
|
|
base_mod.BaseAnswerResource(),
|
|
[
|
|
{
|
|
"call_id": "a1",
|
|
"name": "delete_thing",
|
|
"tool_name": "api_tool",
|
|
"action_name": "delete_thing",
|
|
"pause_type": "awaiting_approval",
|
|
},
|
|
{
|
|
"call_id": "c1",
|
|
"name": "create_file",
|
|
"tool_name": "create_file",
|
|
"action_name": "create_file",
|
|
"pause_type": "requires_client_execution",
|
|
},
|
|
],
|
|
)
|
|
published.assert_called_once()
|
|
args, _ = published.call_args
|
|
assert args[1] == "tool.approval.required"
|
|
summaries = args[2]["pending_tool_calls"]
|
|
assert [s["call_id"] for s in summaries] == ["a1"]
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestProcessResponseStream:
|
|
pass
|
|
|
|
def test_processes_complete_stream(self, mock_mongo_db, flask_app):
|
|
import json
|
|
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
|
|
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',
|
|
f'data: {json.dumps({"type": "source", "source": [{"title": "doc1"}]})}\n\n',
|
|
f'data: {json.dumps({"type": "id", "id": conv_id})}\n\n',
|
|
f'data: {json.dumps({"type": "end"})}\n\n',
|
|
]
|
|
|
|
result = resource.process_response_stream(iter(stream))
|
|
|
|
assert result["conversation_id"] == conv_id
|
|
assert result["answer"] == "Hello world"
|
|
assert result["sources"] == [{"title": "doc1"}]
|
|
assert result["error"] is None
|
|
|
|
def test_handles_stream_error(self, mock_mongo_db, flask_app):
|
|
import json
|
|
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
|
|
stream = [
|
|
f'data: {json.dumps({"type": "error", "error": "Test error"})}\n\n',
|
|
]
|
|
|
|
result = resource.process_response_stream(iter(stream))
|
|
|
|
assert result["conversation_id"] is None
|
|
assert result["error"] == "Test error"
|
|
|
|
def test_handles_malformed_stream_data(self, mock_mongo_db, flask_app):
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
|
|
stream = [
|
|
"data: invalid json\n\n",
|
|
'data: {"type": "end"}\n\n',
|
|
]
|
|
|
|
result = resource.process_response_stream(iter(stream))
|
|
|
|
assert result is not None
|
|
|
|
|
|
@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
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
|
|
error_stream = list(resource.error_stream_generate("Test error message"))
|
|
|
|
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
|