feat(B.13-B.17+B.16): Error-Handling-Infra, Observability, Graceful Shutdown, Cost-Cap, API Versioning
B.13 Error-Handling-Infrastruktur:
- ErrorCategory Enum (TRANSIENT/PERMANENT/PARTIAL), ApiError erweitert
- Einheitliches Error-Response-Format: {code, detail, field, trace_id, retryable, category}
- 3 FastAPI Exception-Handler (ApiError, HTTPException, unhandled)
- classify_exception() Helper, 6 neue Error-Codes
- 28 Tests in test_error_handling.py
B.14 Observability & trace_id-Korrelation:
- trace_id pro Request (UUID4 short) in structlog contextvars
- X-Trace-Id Response-Header
- Sensitive Fields structlog processor
- llm_complete()/llm_embed() akzeptieren trace_id kwarg
- 12 Tests in test_observability.py
B.15 Graceful Shutdown & Connection Draining:
- _shutdown_event + _inflight_requests Tracking in main.py
- drain_all_connections() in ws_helpers.py
- Worker on_shutdown pausiert WorkflowInstances (status=paused)
- 8 Tests in test_graceful_shutdown.py
B.16 API Versioning Strategie:
- Plugin-Dev-Guide Kapitel 30: URL-basiertes Versioning, Breaking Change Prozess
B.17 Cost Overrun Protection:
- llm_monthly_budget_usd + llm_hard_cutoff Settings
- _check_tenant_budget() vor jedem LLM-Call
- _track_tenant_cost() in Redis (INCRBYFLOAT)
- _check_cost_alerts() bei 50%/80%/100% -> post_system_message()
- 20 Tests in test_cost_protection.py
Total: 68 neue Tests, alle grün. Keine Regressionen.
This commit is contained in:
@@ -0,0 +1,144 @@
|
||||
"""Tests for graceful shutdown & connection draining (B.15).
|
||||
|
||||
Covers:
|
||||
- SIGTERM → shutdown event set
|
||||
- WS drain → all connections closed
|
||||
- Worker shutdown → on_shutdown called
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from starlette.websockets import WebSocket
|
||||
|
||||
from app.core.ws_helpers import (
|
||||
drain_all_connections,
|
||||
register_ws_registry,
|
||||
_global_ws_registries,
|
||||
)
|
||||
|
||||
|
||||
class TestShutdownEvent:
|
||||
"""Tests for the API graceful shutdown event."""
|
||||
|
||||
def test_shutdown_event_is_settable(self):
|
||||
"""The _shutdown_event should be settable."""
|
||||
from app.main import _shutdown_event
|
||||
# Reset to clean state
|
||||
_shutdown_event.clear()
|
||||
assert not _shutdown_event.is_set()
|
||||
_shutdown_event.set()
|
||||
assert _shutdown_event.is_set()
|
||||
# Clean up
|
||||
_shutdown_event.clear()
|
||||
|
||||
def test_shutdown_event_is_asyncio_event(self):
|
||||
"""_shutdown_event should be an asyncio.Event."""
|
||||
from app.main import _shutdown_event
|
||||
assert isinstance(_shutdown_event, asyncio.Event)
|
||||
|
||||
|
||||
class TestWebSocketDrain:
|
||||
"""Tests for WebSocket connection draining."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_drain_all_connections_closes_websockets(self):
|
||||
"""drain_all_connections should close all registered WS connections."""
|
||||
# Create mock WebSocket objects
|
||||
ws1 = MagicMock(spec=WebSocket)
|
||||
ws1.send_text = AsyncMock()
|
||||
ws1.close = AsyncMock()
|
||||
ws2 = MagicMock(spec=WebSocket)
|
||||
ws2.send_text = AsyncMock()
|
||||
ws2.close = AsyncMock()
|
||||
|
||||
# Register a test registry
|
||||
test_registry: dict[str, list[WebSocket]] = {
|
||||
"user1": [ws1],
|
||||
"user2": [ws2],
|
||||
}
|
||||
register_ws_registry(test_registry)
|
||||
|
||||
# Drain with 0 grace period for fast test
|
||||
await drain_all_connections(grace_period_seconds=0)
|
||||
|
||||
# Both WebSockets should have received reconnect message and been closed
|
||||
ws1.send_text.assert_called_once()
|
||||
ws2.send_text.assert_called_once()
|
||||
ws1.close.assert_called_once()
|
||||
ws2.close.assert_called_once()
|
||||
|
||||
# Verify reconnect message content
|
||||
call_args = ws1.send_text.call_args[0][0]
|
||||
msg = json.loads(call_args)
|
||||
assert msg["type"] == "reconnect"
|
||||
assert msg["reason"] == "server_shutdown"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_drain_all_connections_handles_empty_registry(self):
|
||||
"""drain_all_connections should handle empty registries gracefully."""
|
||||
# Clear any existing registries from previous tests
|
||||
_global_ws_registries.clear()
|
||||
await drain_all_connections(grace_period_seconds=0)
|
||||
# Should not raise
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_drain_all_connections_handles_send_failure(self):
|
||||
"""drain_all_connections should not fail if send_text raises."""
|
||||
ws = MagicMock(spec=WebSocket)
|
||||
ws.send_text = AsyncMock(side_effect=Exception("WS closed"))
|
||||
ws.close = AsyncMock(side_effect=Exception("WS closed"))
|
||||
|
||||
test_registry: dict[str, list[WebSocket]] = {"user1": [ws]}
|
||||
register_ws_registry(test_registry)
|
||||
|
||||
# Should not raise despite send failures
|
||||
await drain_all_connections(grace_period_seconds=0)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_register_ws_registry_dedup(self):
|
||||
"""register_ws_registry should not add the same registry twice."""
|
||||
_global_ws_registries.clear()
|
||||
reg: dict[str, list[WebSocket]] = {}
|
||||
register_ws_registry(reg)
|
||||
register_ws_registry(reg)
|
||||
assert len(_global_ws_registries) == 1
|
||||
|
||||
|
||||
class TestWorkerShutdown:
|
||||
"""Tests for ARQ worker graceful shutdown."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_on_shutdown_called(self):
|
||||
"""on_shutdown should be callable and close Redis."""
|
||||
from app.core.worker import on_shutdown
|
||||
|
||||
with patch("app.core.auth.close_redis", new_callable=AsyncMock) as mock_close:
|
||||
with patch("app.core.db.get_worker_session_factory") as mock_factory:
|
||||
# Mock the session factory to avoid DB access
|
||||
mock_session = AsyncMock()
|
||||
mock_result = MagicMock()
|
||||
mock_result.scalars.return_value.all.return_value = []
|
||||
mock_session.execute = AsyncMock(return_value=mock_result)
|
||||
mock_session.commit = AsyncMock()
|
||||
mock_ctx = AsyncMock()
|
||||
mock_ctx.__aenter__ = AsyncMock(return_value=mock_session)
|
||||
mock_ctx.__aexit__ = AsyncMock(return_value=False)
|
||||
mock_factory.return_value = mock_ctx
|
||||
|
||||
await on_shutdown({})
|
||||
mock_close.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_on_shutdown_handles_db_error(self):
|
||||
"""on_shutdown should handle DB errors gracefully and still close Redis."""
|
||||
from app.core.worker import on_shutdown
|
||||
|
||||
with patch("app.core.auth.close_redis", new_callable=AsyncMock) as mock_close:
|
||||
with patch("app.core.db.get_worker_session_factory", side_effect=Exception("DB unavailable")):
|
||||
await on_shutdown({})
|
||||
mock_close.assert_called_once()
|
||||
Reference in New Issue
Block a user