wangweiandCopilot
49ee50c104
Critical: the MCP SDK auto-enables DNS-rebinding protection when its host parameter is left at the 127.0.0.1 default, hard-coding a loopback-only Host allow-list. Every remote client (the only deployment this feature targets) was refused with HTTP 421 before auth or the tool ran. Now driven by a new MCP_ALLOWED_HOSTS setting, with '*' as an explicit, logged opt-out. Also bounds query/top_k to match AskRequest (top_k is amplified 4x downstream, so an unbounded value was a resource-exhaustion vector), decodes the Authorization header as latin-1 per the ASGI spec instead of raising a 500 on malformed bytes, and returns WWW-Authenticate on 401 per RFC 7235. Moves the psycopg2 import guard into backend/tests/conftest.py: duplicated across four test modules, it only worked because of alphabetical collection order, and any earlier-sorting package would have reintroduced a live connection attempt against the production database. Registers the mcp module in the authoritative backend architecture doc. 84 backend tests pass. Verified against a live server: allowed remote Host returns a valid initialize result, unknown Host returns 421, missing token returns 401 with WWW-Authenticate. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
103 lines
3.9 KiB
Python
103 lines
3.9 KiB
Python
"""Unit tests for PostgresModelUsageStore, using a mocked psycopg2 pool.
|
|
|
|
Mirrors the mocking pattern in backend/tests/perception/test_postgres_event_store.py
|
|
— no real database is needed.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from datetime import datetime, timezone
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
# psycopg2 is mocked centrally in backend/tests/conftest.py, so importing the
|
|
# module under test here never binds the real driver.
|
|
from app.shared.model_usage_tracker import ModelUsageEntry
|
|
|
|
|
|
def _cursor_returning(rows):
|
|
"""Build a MagicMock standing in for a psycopg2 cursor context manager."""
|
|
cursor = MagicMock()
|
|
cursor.__enter__ = lambda s: s
|
|
cursor.__exit__ = MagicMock(return_value=False)
|
|
cursor.fetchall.return_value = rows
|
|
return cursor
|
|
|
|
|
|
@patch("app.infrastructure.storage.postgres_model_usage_store.PostgresModelUsageStore._ensure_schema")
|
|
@patch("app.infrastructure.storage.postgres_model_usage_store.ThreadedConnectionPool")
|
|
def test_load_all_returns_entries_keyed_by_provider_model(mock_pool_class, mock_ensure):
|
|
"""load_all() must turn each row into a ModelUsageEntry keyed by 'provider:model'."""
|
|
row = {
|
|
"provider": "deepseek",
|
|
"model": "deepseek-v4-flash",
|
|
"total_tokens": 100,
|
|
"prompt_tokens": 60,
|
|
"completion_tokens": 40,
|
|
"call_count_ok": 5,
|
|
"call_count_error": 1,
|
|
"last_called_at": datetime(2026, 7, 23, tzinfo=timezone.utc),
|
|
"last_latency_ms": 250,
|
|
"last_error": None,
|
|
}
|
|
mock_pool = MagicMock()
|
|
mock_pool_class.return_value = mock_pool
|
|
conn = MagicMock()
|
|
conn.__enter__ = lambda s: s
|
|
conn.__exit__ = MagicMock(return_value=False)
|
|
conn.cursor.return_value = _cursor_returning([row])
|
|
mock_pool.getconn.return_value = conn
|
|
|
|
from app.infrastructure.storage.postgres_model_usage_store import PostgresModelUsageStore
|
|
store = PostgresModelUsageStore()
|
|
entries = store.load_all()
|
|
|
|
assert "deepseek:deepseek-v4-flash" in entries
|
|
entry = entries["deepseek:deepseek-v4-flash"]
|
|
assert isinstance(entry, ModelUsageEntry)
|
|
assert entry.total_tokens == 100
|
|
assert entry.call_count_error == 1
|
|
|
|
|
|
@patch("app.infrastructure.storage.postgres_model_usage_store.PostgresModelUsageStore._ensure_schema")
|
|
@patch("app.infrastructure.storage.postgres_model_usage_store.ThreadedConnectionPool")
|
|
def test_flush_upserts_every_entry(mock_pool_class, mock_ensure):
|
|
"""flush() must execute one UPSERT per tracked entry and commit once."""
|
|
mock_pool = MagicMock()
|
|
mock_pool_class.return_value = mock_pool
|
|
conn = MagicMock()
|
|
conn.__enter__ = lambda s: s
|
|
conn.__exit__ = MagicMock(return_value=False)
|
|
cursor = MagicMock()
|
|
cursor.__enter__ = lambda s: s
|
|
cursor.__exit__ = MagicMock(return_value=False)
|
|
conn.cursor.return_value = cursor
|
|
mock_pool.getconn.return_value = conn
|
|
|
|
from app.infrastructure.storage.postgres_model_usage_store import PostgresModelUsageStore
|
|
store = PostgresModelUsageStore()
|
|
entries = {
|
|
"deepseek:deepseek-v4-flash": ModelUsageEntry(
|
|
provider="deepseek", model="deepseek-v4-flash", total_tokens=100, call_count_ok=5,
|
|
),
|
|
}
|
|
|
|
store.flush(entries)
|
|
|
|
assert cursor.execute.call_count == 1
|
|
conn.commit.assert_called_once()
|
|
|
|
|
|
@patch("app.infrastructure.storage.postgres_model_usage_store.PostgresModelUsageStore._ensure_schema")
|
|
@patch("app.infrastructure.storage.postgres_model_usage_store.ThreadedConnectionPool")
|
|
def test_flush_with_no_entries_does_not_touch_the_database(mock_pool_class, mock_ensure):
|
|
"""flush({}) must be a no-op — no point opening a connection for nothing."""
|
|
mock_pool = MagicMock()
|
|
mock_pool_class.return_value = mock_pool
|
|
|
|
from app.infrastructure.storage.postgres_model_usage_store import PostgresModelUsageStore
|
|
store = PostgresModelUsageStore()
|
|
|
|
store.flush({})
|
|
|
|
mock_pool.getconn.assert_not_called()
|