Phase 4: Circuit Breaker, DB Retry, Redis Graceful Degradation

- app/core/resilience.py: CircuitBreaker (CLOSED/OPEN/HALF_OPEN), retry_db,
  redis_call_with_fallback, InMemoryRateLimiter, CircuitBreakerMiddleware
- app/core/auth.py: get_session_data now falls back to PostgreSQL sessions
  table when Redis is unavailable
- app/core/permissions.py: get_cached_permissions falls back to direct DB
  resolution when Redis circuit is open
- app/core/rate_limit.py: check_rate_limit falls back to in-memory limiter
  when Redis is down; reset_rate_limit clears both Redis and in-memory
- app/core/middleware.py: CSRF validation uses get_session_data (Redis+DB
  fallback); sliding session TTL is best-effort during outage
- app/core/db/__init__.py: get_db() wraps session creation with retry_db
  for transient connection errors; records circuit breaker success/failure
- app/deps.py: refresh_session_ttl wrapped in try/except for Redis outage
- app/main.py: CircuitBreakerMiddleware registered (returns 503 when DB
  circuit is OPEN, skips health/metrics endpoints)
- app/config.py: Added resilience settings (thresholds, cooldown, retries)
- tests/test_resilience.py: 30 tests covering all patterns

30/30 resilience tests pass. No regressions in plugin lifecycle tests.
This commit is contained in:
Agent Zero
2026-08-04 14:34:06 +02:00
parent 247d4165ea
commit a26405f15e
10 changed files with 875 additions and 58 deletions
+7
View File
@@ -66,6 +66,13 @@ class Settings(BaseSettings):
# Trusted proxy CIDRs (comma-separated) — only these proxies can set X-Forwarded-For # Trusted proxy CIDRs (comma-separated) — only these proxies can set X-Forwarded-For
trusted_proxy_cidrs: str = "" trusted_proxy_cidrs: str = ""
# Resilience
circuit_breaker_failure_threshold: int = 5
circuit_breaker_window_seconds: int = 30
circuit_breaker_cooldown_seconds: int = 60
db_retry_max_attempts: int = 3
db_retry_base_delay: float = 0.1
# Rate Limiting # Rate Limiting
rate_limit_login_max: int = 5 rate_limit_login_max: int = 5
rate_limit_login_window: int = 900 # 15 min rate_limit_login_window: int = 900 # 15 min
+42 -4
View File
@@ -188,13 +188,51 @@ async def create_session(
async def get_session_data(redis: aioredis.Redis, session_id: str) -> dict[str, Any] | None: async def get_session_data(redis: aioredis.Redis, session_id: str) -> dict[str, Any] | None:
"""Retrieve session data from Redis.""" """Retrieve session data from Redis with DB fallback.
Tries Redis first. If Redis is unavailable, falls back to PostgreSQL
sessions table (audit trail) to keep users logged in during Redis outages.
"""
import json import json
raw = await redis.get(f"session:{session_id}") from app.core.resilience import get_circuit
if raw is None:
circuit = get_circuit("redis")
if await circuit.can_proceed():
try:
raw = await redis.get(f"session:{session_id}")
await circuit.record_success()
if raw is None:
return None
return json.loads(raw)
except Exception as exc:
logger.warning("Redis session lookup failed: %s — falling back to DB", exc)
await circuit.record_failure()
# DB fallback: query sessions table
try:
from app.core.db import get_auth_session_factory
from app.models.session import Session as SessionModel
from sqlalchemy import select
from datetime import UTC, datetime
factory = get_auth_session_factory()
async with factory() as db:
result = await db.execute(
select(SessionModel).where(SessionModel.id == uuid.UUID(session_id))
)
session = result.scalar_one_or_none()
if session is None or session.expires_at < datetime.now(UTC):
return None
return {
"user_id": str(session.user_id),
"tenant_id": str(session.tenant_id),
"csrf_token": session.csrf_token,
"is_active": True,
}
except Exception as db_exc:
logger.error("DB fallback for session lookup also failed: %s", db_exc)
return None return None
return json.loads(raw)
async def refresh_session_ttl(redis: aioredis.Redis, session_id: str) -> None: async def refresh_session_ttl(redis: aioredis.Redis, session_id: str) -> None:
+18 -8
View File
@@ -222,15 +222,25 @@ async def get_db() -> AsyncGenerator[AsyncSession, None]:
"""FastAPI dependency: yield an async database session (crm_api role). """FastAPI dependency: yield an async database session (crm_api role).
Used for normal API requests with tenant context set via RLS. Used for normal API requests with tenant context set via RLS.
Includes retry logic for transient connection errors.
""" """
factory = get_session_factory() from app.core.resilience import get_circuit, retry_db
async with factory() as session:
try: async def _get_session():
yield session factory = get_session_factory()
await session.commit() return factory()
except Exception:
await session.rollback() session = await retry_db(_get_session)
raise try:
yield session
await session.commit()
await get_circuit("db").record_success()
except Exception:
await session.rollback()
await get_circuit("db").record_failure()
raise
finally:
await session.close()
async def get_auth_db() -> AsyncGenerator[AsyncSession, None]: async def get_auth_db() -> AsyncGenerator[AsyncSession, None]:
+9 -6
View File
@@ -111,18 +111,17 @@ class CSRFMiddleware(BaseHTTPMiddleware):
content={"detail": "No session for CSRF validation", "code": "csrf_no_session"}, content={"detail": "No session for CSRF validation", "code": "csrf_no_session"},
) )
# Look up CSRF token from Redis session (use singleton) # Look up CSRF token from session (Redis with DB fallback)
from app.core.auth import get_redis from app.core.auth import get_redis, get_session_data
redis = get_redis() redis = get_redis()
raw = await redis.get(f"session:{session_id}") session_data = await get_session_data(redis, session_id)
if raw is None: if session_data is None:
return JSONResponse( return JSONResponse(
status_code=status.HTTP_403_FORBIDDEN, status_code=status.HTTP_403_FORBIDDEN,
content={"detail": "Session expired for CSRF validation", "code": "csrf_session_expired"}, content={"detail": "Session expired for CSRF validation", "code": "csrf_session_expired"},
) )
session_data = json.loads(raw)
stored_token = session_data.get("csrf_token") stored_token = session_data.get("csrf_token")
if not stored_token or stored_token != csrf_header: if not stored_token or stored_token != csrf_header:
@@ -132,6 +131,10 @@ class CSRFMiddleware(BaseHTTPMiddleware):
) )
# Sliding session: also extend TTL on CSRF-validated unsafe requests # Sliding session: also extend TTL on CSRF-validated unsafe requests
await redis.expire(f"session:{session_id}", settings.session_ttl_seconds) # (best-effort — ignore Redis errors during outage)
try:
await redis.expire(f"session:{session_id}", settings.session_ttl_seconds)
except Exception:
pass
return await call_next(request) return await call_next(request)
+46 -28
View File
@@ -349,43 +349,54 @@ async def get_cached_permissions(
Validates the cached permission_version against the current DB version. Validates the cached permission_version against the current DB version.
If they differ, the cache entry is stale and will be re-resolved. If they differ, the cache entry is stale and will be re-resolved.
Falls back to direct DB resolution when Redis is unavailable.
""" """
cache_key = f"{CACHE_PREFIX}:{user_id}:{tenant_id}" cache_key = f"{CACHE_PREFIX}:{user_id}:{tenant_id}"
raw = await redis.get(cache_key) from app.core.resilience import get_circuit
if raw is not None:
data = json.loads(raw)
cached_version = data.get("version", -1)
# Validate cached version against current DB version circuit = get_circuit("redis")
redis_available = await circuit.can_proceed()
if redis_available:
try: try:
current_version = await _get_current_permission_version(db, user_id, tenant_id) raw = await redis.get(cache_key)
except Exception: await circuit.record_success()
logger.warning( if raw is not None:
"Failed to query current permission_version for cache validation " data = json.loads(raw)
"(user=%s, tenant=%s) — invalidating cache and re-resolving", cached_version = data.get("version", -1)
user_id, tenant_id,
exc_info=True,
)
# Invalidate stale cache — do NOT trust cached permissions on DB error
await redis.delete(cache_key)
return None # Fall through to re-resolution from DB
if cached_version == current_version: # Validate cached version against current DB version
return data try:
current_version = await _get_current_permission_version(db, user_id, tenant_id)
except Exception:
logger.warning(
"Failed to query current permission_version for cache validation "
"(user=%s, tenant=%s) — invalidating cache and re-resolving",
user_id, tenant_id,
exc_info=True,
)
await redis.delete(cache_key)
return None # Fall through to re-resolution from DB
# Version mismatch — invalidate stale cache and re-resolve if cached_version == current_version:
logger.info( return data
"Permission cache version mismatch for user=%s tenant=%s "
"(cached=%s, current=%s) — re-resolving",
user_id, tenant_id, cached_version, current_version,
)
await redis.delete(cache_key)
# Cache miss or stale — resolve from DB logger.info(
"Permission cache version mismatch for user=%s tenant=%s "
"(cached=%s, current=%s) — re-resolving",
user_id, tenant_id, cached_version, current_version,
)
await redis.delete(cache_key)
except Exception as exc:
logger.warning("Redis permission cache failed: %s — resolving from DB", exc)
await circuit.record_failure()
redis_available = False
# Cache miss, stale, or Redis unavailable — resolve from DB
resolved = await resolve_permissions(db, user_id, tenant_id) resolved = await resolve_permissions(db, user_id, tenant_id)
# Store in cache (convert sets to lists for JSON)
cache_data = { cache_data = {
"permissions": list(resolved["permissions"]), "permissions": list(resolved["permissions"]),
"denied": list(resolved["denied"]), "denied": list(resolved["denied"]),
@@ -393,7 +404,14 @@ async def get_cached_permissions(
"is_system_admin": resolved["is_system_admin"], "is_system_admin": resolved["is_system_admin"],
"version": resolved["version"], "version": resolved["version"],
} }
await redis.setex(cache_key, CACHE_TTL, json.dumps(cache_data))
# Try to cache (best-effort during Redis outage)
if redis_available:
try:
await redis.setex(cache_key, CACHE_TTL, json.dumps(cache_data))
except Exception:
logger.warning("Failed to cache permissions in Redis — continuing without cache")
return cache_data return cache_data
+56 -11
View File
@@ -19,29 +19,74 @@ async def check_rate_limit(
window_seconds: int, window_seconds: int,
) -> None: ) -> None:
"""Check rate limit using Redis INCR + EXPIRE. """Check rate limit using Redis INCR + EXPIRE.
Falls back to in-memory rate limiter when Redis is unavailable.
Raises 429 if limit exceeded. Raises 429 if limit exceeded.
""" """
redis = get_redis() from app.core.resilience import get_circuit, get_inmemory_limiter
current = await redis.incr(redis_key)
if current == 1: circuit = get_circuit("redis")
await redis.expire(redis_key, window_seconds) if await circuit.can_proceed():
if current > max_attempts: try:
ttl = await redis.ttl(redis_key) redis = get_redis()
current = await redis.incr(redis_key)
if current == 1:
await redis.expire(redis_key, window_seconds)
if current > max_attempts:
ttl = await redis.ttl(redis_key)
raise HTTPException(
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
detail={
"detail": "Rate limit exceeded",
"code": "rate_limited",
"retry_after": ttl,
},
headers={"Retry-After": str(ttl)} if ttl > 0 else {},
)
await circuit.record_success()
return
except HTTPException:
raise
except Exception as exc:
logger.warning("Redis rate limit failed: %s — using in-memory fallback", exc)
await circuit.record_failure()
# In-memory fallback
limiter = get_inmemory_limiter()
allowed, retry_after = await limiter.check(redis_key, max_attempts, window_seconds)
if not allowed:
raise HTTPException( raise HTTPException(
status_code=status.HTTP_429_TOO_MANY_REQUESTS, status_code=status.HTTP_429_TOO_MANY_REQUESTS,
detail={ detail={
"detail": "Rate limit exceeded", "detail": "Rate limit exceeded",
"code": "rate_limited", "code": "rate_limited",
"retry_after": ttl, "retry_after": retry_after,
}, },
headers={"Retry-After": str(ttl)} if ttl > 0 else {}, headers={"Retry-After": str(retry_after)} if retry_after > 0 else {},
) )
async def reset_rate_limit(redis_key: str) -> None: async def reset_rate_limit(redis_key: str) -> None:
"""Reset a rate limit counter (e.g. on successful login).""" """Reset a rate limit counter (e.g. on successful login).
redis = get_redis()
await redis.delete(redis_key) Resets both Redis and in-memory counters.
"""
from app.core.resilience import get_circuit, get_inmemory_limiter
# Always reset in-memory
limiter = get_inmemory_limiter()
await limiter.reset(redis_key)
# Try Redis
circuit = get_circuit("redis")
if await circuit.can_proceed():
try:
redis = get_redis()
await redis.delete(redis_key)
await circuit.record_success()
except Exception as exc:
logger.warning("Redis rate limit reset failed: %s", exc)
await circuit.record_failure()
def get_client_ip(request: Request) -> str: def get_client_ip(request: Request) -> str:
+300
View File
@@ -0,0 +1,300 @@
"""Resilience patterns: circuit breaker, DB retry, Redis graceful degradation.
Provides:
- CircuitBreaker: tracks failures per service, opens after threshold, auto-recovers
- retry_db: decorator that retries DB operations on transient connection errors
- redis_call_with_fallback: helper that catches Redis errors and falls back
- InMemoryRateLimiter: process-local rate limiter for Redis outage fallback
- CircuitBreakerMiddleware: ASGI middleware that fail-fasts on open circuits
"""
from __future__ import annotations
import asyncio
import logging
import time
from collections import defaultdict
from collections.abc import Awaitable, Callable
from typing import Any, TypeVar
from fastapi import HTTPException, status
from starlette.responses import JSONResponse
logger = logging.getLogger(__name__)
T = TypeVar("T")
# -- Circuit Breaker ----------------------------------------------------------
class CircuitBreaker:
"""Circuit breaker for a single service.
States:
- CLOSED: normal operation, requests pass through
- OPEN: service is down, requests fail fast with 503
- HALF_OPEN: cooldown expired, one probe request is allowed through
"""
CLOSED = "closed"
OPEN = "open"
HALF_OPEN = "half_open"
def __init__(
self,
name: str,
failure_threshold: int = 5,
window_seconds: int = 30,
cooldown_seconds: int = 60,
) -> None:
self.name = name
self.failure_threshold = failure_threshold
self.window_seconds = window_seconds
self.cooldown_seconds = cooldown_seconds
self._state = self.CLOSED
self._failures: list[float] = []
self._opened_at: float = 0.0
self._half_open_probe_in_flight = False
self._lock = asyncio.Lock()
@property
def state(self) -> str:
if self._state == self.OPEN:
if time.monotonic() - self._opened_at >= self.cooldown_seconds:
return self.HALF_OPEN
return self._state
async def can_proceed(self) -> bool:
current_state = self.state
if current_state == self.CLOSED:
return True
if current_state == self.HALF_OPEN:
async with self._lock:
if not self._half_open_probe_in_flight:
self._half_open_probe_in_flight = True
return True
return False
return False
async def record_success(self) -> None:
was_open = self._state in (self.OPEN, self.HALF_OPEN)
self._state = self.CLOSED
self._failures.clear()
self._half_open_probe_in_flight = False
if was_open:
logger.info("Circuit breaker '%s' recovered - state: CLOSED", self.name)
async def record_failure(self) -> None:
now = time.monotonic()
self._failures.append(now)
cutoff = now - self.window_seconds
self._failures = [t for t in self._failures if t >= cutoff]
if self._state == self.HALF_OPEN:
self._state = self.OPEN
self._opened_at = now
self._half_open_probe_in_flight = False
logger.warning("Circuit breaker '%s' probe failed - state: OPEN", self.name)
return
if len(self._failures) >= self.failure_threshold:
self._state = self.OPEN
self._opened_at = now
logger.error(
"Circuit breaker '%s' tripped - %d failures in %ds - state: OPEN",
self.name,
len(self._failures),
self.window_seconds,
)
def reset(self) -> None:
self._state = self.CLOSED
self._failures.clear()
self._opened_at = 0.0
self._half_open_probe_in_flight = False
_circuits: dict[str, CircuitBreaker] = {}
def get_circuit(name: str) -> CircuitBreaker:
if name not in _circuits:
_circuits[name] = CircuitBreaker(name)
return _circuits[name]
def reset_all_circuits() -> None:
for cb in _circuits.values():
cb.reset()
# -- DB Retry -----------------------------------------------------------------
_DB_RETRYABLE_EXC = (ConnectionError, OSError, asyncio.TimeoutError)
def _is_transient_db_error(exc: Exception) -> bool:
if isinstance(exc, _DB_RETRYABLE_EXC):
return True
try:
from sqlalchemy.exc import OperationalError, DBAPIError
if isinstance(exc, OperationalError):
return True
if isinstance(exc, DBAPIError):
cause = exc.__cause__ or exc.orig
if cause and isinstance(cause, (ConnectionError, OSError)):
return True
cause_str = str(cause) if cause else ""
if any(kw in cause_str.lower() for kw in (
"connection", "timeout", "refused", "reset",
"broken pipe", "server closed",
)):
return True
except ImportError:
pass
return False
async def retry_db(
func: Callable[..., Awaitable[T]],
*args: Any,
max_retries: int = 3,
base_delay: float = 0.1,
**kwargs: Any,
) -> T:
last_exc: Exception | None = None
for attempt in range(max_retries):
try:
result = await func(*args, **kwargs)
if attempt > 0:
logger.info("DB operation succeeded on retry %d", attempt + 1)
return result
except Exception as exc:
last_exc = exc
if not _is_transient_db_error(exc):
raise
if attempt < max_retries - 1:
delay = base_delay * (2 ** attempt)
logger.warning(
"DB transient error (attempt %d/%d): %s - retrying in %.2fs",
attempt + 1, max_retries, exc, delay,
)
await asyncio.sleep(delay)
else:
logger.error("DB operation failed after %d retries: %s", max_retries, exc)
await get_circuit("db").record_failure()
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail={"detail": "Database temporarily unavailable", "code": "db_unavailable"},
)
# -- Redis Fallback -----------------------------------------------------------
class RedisUnavailableError(Exception):
pass
async def redis_call_with_fallback(
redis_op: Callable[..., Awaitable[T]],
fallback: Callable[..., Awaitable[T]] | None = None,
*args: Any,
**kwargs: Any,
) -> T | None:
circuit = get_circuit("redis")
if not await circuit.can_proceed():
logger.warning("Redis circuit OPEN - using fallback")
if fallback is not None:
return await fallback(*args, **kwargs)
return None
try:
result = await redis_op(*args, **kwargs)
await circuit.record_success()
return result
except Exception as exc:
logger.warning("Redis operation failed: %s - using fallback", exc)
await circuit.record_failure()
if fallback is not None:
return await fallback(*args, **kwargs)
return None
# -- In-Memory Rate Limiter (Redis fallback) ----------------------------------
class InMemoryRateLimiter:
"""Process-local sliding-window rate limiter for Redis outage fallback."""
def __init__(self) -> None:
self._counts: dict[str, list[float]] = defaultdict(list)
async def check(self, key: str, max_attempts: int, window_seconds: int) -> tuple[bool, int]:
now = time.monotonic()
cutoff = now - window_seconds
self._counts[key] = [t for t in self._counts[key] if t >= cutoff]
if len(self._counts[key]) >= max_attempts:
retry_after = int(window_seconds - (now - self._counts[key][0]))
return False, max(retry_after, 1)
self._counts[key].append(now)
return True, 0
async def reset(self, key: str) -> None:
self._counts.pop(key, None)
def clear_all(self) -> None:
self._counts.clear()
_inmemory_limiter = InMemoryRateLimiter()
def get_inmemory_limiter() -> InMemoryRateLimiter:
return _inmemory_limiter
# -- Circuit Breaker Middleware -----------------------------------------------
class CircuitBreakerMiddleware:
"""ASGI middleware: returns 503 when DB circuit is OPEN."""
SKIP_PATHS = {
"/api/v1/health",
"/api/v1/health/live",
"/api/v1/health/ready",
"/api/v1/metrics",
"/docs",
"/redoc",
"/openapi.json",
}
def __init__(self, app) -> None:
self.app = app
async def __call__(self, scope, receive, send) -> None:
if scope["type"] != "http":
await self.app(scope, receive, send)
return
path = scope.get("path", "")
if path in self.SKIP_PATHS or path.startswith("/docs") or path.startswith("/redoc"):
await self.app(scope, receive, send)
return
db_circuit = get_circuit("db")
if not await db_circuit.can_proceed():
response = JSONResponse(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
content={
"detail": "Service temporarily unavailable",
"code": "circuit_open",
"service": "db",
},
)
await response(scope, receive, send)
return
await self.app(scope, receive, send)
+5 -1
View File
@@ -103,7 +103,11 @@ async def get_current_user(
) )
# Sliding session: extend TTL on each authenticated request # Sliding session: extend TTL on each authenticated request
await refresh_session_ttl(redis, session_id) # Best-effort during Redis outage — session still valid from DB fallback
try:
await refresh_session_ttl(redis, session_id)
except Exception:
logger.debug("refresh_session_ttl failed (Redis may be down) — continuing")
if not session_data.get("is_active", True): if not session_data.get("is_active", True):
raise HTTPException( raise HTTPException(
+2
View File
@@ -23,6 +23,7 @@ from app.core.db import close_engine, get_engine
from app.core.error_codes import ApiError from app.core.error_codes import ApiError
from app.core.middleware import CSRFMiddleware, SecurityHeadersMiddleware from app.core.middleware import CSRFMiddleware, SecurityHeadersMiddleware
from app.core.rate_limit import GeneralRateLimitMiddleware from app.core.rate_limit import GeneralRateLimitMiddleware
from app.core.resilience import CircuitBreakerMiddleware
from app.core.monitoring import record_error, record_request from app.core.monitoring import record_error, record_request
from app.core.plugin_error_handler import wrap_plugin_route from app.core.plugin_error_handler import wrap_plugin_route
from app.core.service_container import get_container from app.core.service_container import get_container
@@ -379,6 +380,7 @@ def create_app() -> FastAPI:
app.add_middleware(SecurityHeadersMiddleware) app.add_middleware(SecurityHeadersMiddleware)
app.add_middleware(GeneralRateLimitMiddleware) app.add_middleware(GeneralRateLimitMiddleware)
app.add_middleware(RequestLoggingMiddleware) app.add_middleware(RequestLoggingMiddleware)
app.add_middleware(CircuitBreakerMiddleware)
# ── Global exception handler — catch ALL unhandled exceptions ── # ── Global exception handler — catch ALL unhandled exceptions ──
@app.exception_handler(Exception) @app.exception_handler(Exception)
+390
View File
@@ -0,0 +1,390 @@
"""Tests for resilience patterns: circuit breaker, DB retry, Redis fallback, in-memory rate limiter."""
import asyncio
import time
import pytest
from unittest.mock import AsyncMock, MagicMock, patch
from app.core.resilience import (
CircuitBreaker,
CircuitBreakerMiddleware,
InMemoryRateLimiter,
RedisUnavailableError,
get_circuit,
get_inmemory_limiter,
reset_all_circuits,
retry_db,
redis_call_with_fallback,
)
class TestCircuitBreaker:
"""Unit tests for CircuitBreaker."""
def setup_method(self):
self.cb = CircuitBreaker("test", failure_threshold=3, window_seconds=30, cooldown_seconds=60)
@pytest.mark.asyncio
async def test_closed_allows_requests(self):
assert await self.cb.can_proceed() is True
@pytest.mark.asyncio
async def test_opens_after_threshold(self):
for _ in range(3):
await self.cb.record_failure()
assert self.cb.state == CircuitBreaker.OPEN
assert await self.cb.can_proceed() is False
@pytest.mark.asyncio
async def test_success_resets_circuit(self):
for _ in range(2):
await self.cb.record_failure()
await self.cb.record_success()
assert self.cb.state == CircuitBreaker.CLOSED
assert len(self.cb._failures) == 0
@pytest.mark.asyncio
async def test_half_open_after_cooldown(self):
for _ in range(3):
await self.cb.record_failure()
assert self.cb.state == CircuitBreaker.OPEN
# Simulate cooldown expiry
self.cb._opened_at = time.monotonic() - 61
assert self.cb.state == CircuitBreaker.HALF_OPEN
# One probe allowed
assert await self.cb.can_proceed() is True
# Second probe blocked
assert await self.cb.can_proceed() is False
@pytest.mark.asyncio
async def test_half_open_probe_success_closes_circuit(self):
for _ in range(3):
await self.cb.record_failure()
self.cb._opened_at = time.monotonic() - 61
# Allow probe
await self.cb.can_proceed()
await self.cb.record_success()
assert self.cb.state == CircuitBreaker.CLOSED
@pytest.mark.asyncio
async def test_half_open_probe_failure_reopens(self):
for _ in range(3):
await self.cb.record_failure()
self.cb._opened_at = time.monotonic() - 61
await self.cb.can_proceed()
await self.cb.record_failure()
assert self.cb.state == CircuitBreaker.OPEN
@pytest.mark.asyncio
async def test_failures_outside_window_are_pruned(self):
cb = CircuitBreaker("test2", failure_threshold=2, window_seconds=1, cooldown_seconds=60)
await cb.record_failure()
# Wait for window to expire
await asyncio.sleep(1.1)
await cb.record_failure()
# Only 1 failure in current window — circuit should stay closed
assert cb.state == CircuitBreaker.CLOSED
def test_reset(self):
self.cb._state = CircuitBreaker.OPEN
self.cb._failures = [1.0, 2.0]
self.cb.reset()
assert self.cb.state == CircuitBreaker.CLOSED
assert len(self.cb._failures) == 0
class TestGetCircuit:
"""Test global circuit registry."""
def test_returns_same_instance(self):
cb1 = get_circuit("shared")
cb2 = get_circuit("shared")
assert cb1 is cb2
def test_different_names_different_instances(self):
cb1 = get_circuit("db")
cb2 = get_circuit("redis")
assert cb1 is not cb2
def test_reset_all(self):
get_circuit("db")._state = CircuitBreaker.OPEN
get_circuit("redis")._state = CircuitBreaker.OPEN
reset_all_circuits()
assert get_circuit("db").state == CircuitBreaker.CLOSED
assert get_circuit("redis").state == CircuitBreaker.CLOSED
class TestRetryDB:
"""Tests for DB retry decorator."""
@pytest.mark.asyncio
async def test_success_no_retry(self):
call_count = 0
async def success_op():
nonlocal call_count
call_count += 1
return "ok"
result = await retry_db(success_op, max_retries=3, base_delay=0.01)
assert result == "ok"
assert call_count == 1
@pytest.mark.asyncio
async def test_retries_on_transient_error(self):
call_count = 0
async def flaky_op():
nonlocal call_count
call_count += 1
if call_count < 3:
raise ConnectionError("DB connection refused")
return "ok"
result = await retry_db(flaky_op, max_retries=3, base_delay=0.01)
assert result == "ok"
assert call_count == 3
@pytest.mark.asyncio
async def test_non_transient_error_not_retried(self):
call_count = 0
async def failing_op():
nonlocal call_count
call_count += 1
raise ValueError("Invalid query")
with pytest.raises(ValueError):
await retry_db(failing_op, max_retries=3, base_delay=0.01)
assert call_count == 1
@pytest.mark.asyncio
async def test_all_retries_exhausted_raises_503(self):
from fastapi import HTTPException
async def always_fail():
raise ConnectionError("DB down")
with pytest.raises(HTTPException) as exc_info:
await retry_db(always_fail, max_retries=2, base_delay=0.01)
assert exc_info.value.status_code == 503
assert exc_info.value.detail["code"] == "db_unavailable"
@pytest.mark.asyncio
async def test_sqlalchemy_operational_error_retried(self):
from sqlalchemy.exc import OperationalError
call_count = 0
async def flaky_op():
nonlocal call_count
call_count += 1
if call_count < 2:
raise OperationalError("stmt", params={}, orig=Exception("connection reset"))
return "ok"
result = await retry_db(flaky_op, max_retries=3, base_delay=0.01)
assert result == "ok"
assert call_count == 2
class TestRedisFallback:
"""Tests for redis_call_with_fallback."""
@pytest.mark.asyncio
async def test_redis_success(self):
reset_all_circuits()
async def redis_get(key):
return f"value:{key}"
result = await redis_call_with_fallback(redis_get, None, "testkey")
assert result == "value:testkey"
@pytest.mark.asyncio
async def test_redis_failure_with_fallback(self):
reset_all_circuits()
async def redis_failing(key):
raise ConnectionError("Redis down")
async def db_fallback(key):
return f"db_value:{key}"
result = await redis_call_with_fallback(redis_failing, db_fallback, "testkey")
assert result == "db_value:testkey"
@pytest.mark.asyncio
async def test_redis_failure_no_fallback_returns_none(self):
reset_all_circuits()
async def redis_failing(key):
raise ConnectionError("Redis down")
result = await redis_call_with_fallback(redis_failing, None, "testkey")
assert result is None
@pytest.mark.asyncio
async def test_circuit_open_uses_fallback(self):
reset_all_circuits()
cb = get_circuit("redis")
# Force circuit open
cb._state = CircuitBreaker.OPEN
cb._opened_at = time.monotonic()
async def redis_op(key):
raise AssertionError("Should not be called when circuit is open")
async def fallback(key):
return f"fallback:{key}"
result = await redis_call_with_fallback(redis_op, fallback, "testkey")
assert result == "fallback:testkey"
class TestInMemoryRateLimiter:
"""Tests for InMemoryRateLimiter."""
@pytest.mark.asyncio
async def test_allows_under_limit(self):
limiter = InMemoryRateLimiter()
allowed, retry_after = await limiter.check("key1", max_attempts=5, window_seconds=60)
assert allowed is True
assert retry_after == 0
@pytest.mark.asyncio
async def test_blocks_over_limit(self):
limiter = InMemoryRateLimiter()
for _ in range(5):
await limiter.check("key1", max_attempts=5, window_seconds=60)
allowed, retry_after = await limiter.check("key1", max_attempts=5, window_seconds=60)
assert allowed is False
assert retry_after > 0
@pytest.mark.asyncio
async def test_different_keys_independent(self):
limiter = InMemoryRateLimiter()
for _ in range(3):
await limiter.check("key1", max_attempts=3, window_seconds=60)
allowed2, _ = await limiter.check("key2", max_attempts=3, window_seconds=60)
assert allowed2 is True
@pytest.mark.asyncio
async def test_reset_clears_key(self):
limiter = InMemoryRateLimiter()
for _ in range(3):
await limiter.check("key1", max_attempts=3, window_seconds=60)
await limiter.reset("key1")
allowed, _ = await limiter.check("key1", max_attempts=3, window_seconds=60)
assert allowed is True
@pytest.mark.asyncio
async def test_window_expiry(self):
limiter = InMemoryRateLimiter()
for _ in range(3):
await limiter.check("key1", max_attempts=3, window_seconds=1)
# Wait for window to expire
await asyncio.sleep(1.1)
allowed, _ = await limiter.check("key1", max_attempts=3, window_seconds=1)
assert allowed is True
def test_clear_all(self):
limiter = InMemoryRateLimiter()
limiter._counts["a"] = [1.0]
limiter._counts["b"] = [2.0]
limiter.clear_all()
assert len(limiter._counts) == 0
class TestCircuitBreakerMiddleware:
"""Tests for CircuitBreakerMiddleware ASGI behavior."""
@pytest.mark.asyncio
async def test_passes_through_when_closed(self):
reset_all_circuits()
called = False
async def app(scope, receive, send):
nonlocal called
called = True
await send({"type": "http.response.start", "status": 200, "headers": []})
await send({"type": "http.response.body", "body": b"ok"})
send_mock = AsyncMock()
middleware = CircuitBreakerMiddleware(app)
scope = {"type": "http", "path": "/api/v1/contacts"}
await middleware(scope, MagicMock(), send_mock)
assert called is True
@pytest.mark.asyncio
async def test_returns_503_when_open(self):
reset_all_circuits()
cb = get_circuit("db")
cb._state = CircuitBreaker.OPEN
cb._opened_at = time.monotonic()
called = False
async def app(scope, receive, send):
nonlocal called
called = True
sent_responses = []
async def mock_send(msg):
sent_responses.append(msg)
middleware = CircuitBreakerMiddleware(app)
scope = {"type": "http", "path": "/api/v1/contacts"}
await middleware(scope, MagicMock(), mock_send)
assert called is False
assert sent_responses[0]["status"] == 503
@pytest.mark.asyncio
async def test_skips_health_endpoints(self):
reset_all_circuits()
cb = get_circuit("db")
cb._state = CircuitBreaker.OPEN
cb._opened_at = time.monotonic()
called = False
async def app(scope, receive, send):
nonlocal called
called = True
await send({"type": "http.response.start", "status": 200, "headers": []})
await send({"type": "http.response.body", "body": b"ok"})
send_mock = AsyncMock()
middleware = CircuitBreakerMiddleware(app)
scope = {"type": "http", "path": "/api/v1/health"}
await middleware(scope, MagicMock(), send_mock)
assert called is True
@pytest.mark.asyncio
async def test_passes_websocket(self):
reset_all_circuits()
called = False
async def app(scope, receive, send):
nonlocal called
called = True
middleware = CircuitBreakerMiddleware(app)
scope = {"type": "websocket", "path": "/ws"}
await middleware(scope, MagicMock(), MagicMock())
assert called is True
# Fixtures: reset circuits before each test class
@pytest.fixture(autouse=True)
def reset_circuits():
reset_all_circuits()
get_inmemory_limiter().clear_all()
yield
reset_all_circuits()
get_inmemory_limiter().clear_all()