Spaces:
Runtime error
Runtime error
File size: 4,223 Bytes
6d20eab 8e12c80 6d20eab 8e12c80 6d20eab 8e12c80 6d20eab 8e12c80 6d20eab 8e12c80 6d20eab 8e12c80 6d20eab | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 | """
pytest configuration and fixtures for ARF API tests.
"""
from app.core.usage_tracker import enforce_quota, Tier
from app.api.deps import get_db, verify_internal_key
from app.database.base import Base
from app.database.models_intents import IntentDB, TenantDB, BetaStateDB, DecisionAuditLogDB # noqa: E501,F401 -- imported for their side effect of registering these tables on Base.metadata
from app.main import app as fastapi_app
from sqlalchemy.orm import sessionmaker
from sqlalchemy import create_engine
from fastapi import Request
from fastapi.testclient import TestClient
import app.core.usage_tracker
import os
import pytest
# ===== STEP 1: Set environment variables BEFORE any app imports =====
os.environ["ARF_USAGE_TRACKING"] = "false"
# Force the correct database URL for tests
os.environ["DATABASE_URL"] = "postgresql://postgres:postgres@localhost:5432/testdb"
os.environ["TEST_DATABASE_URL"] = "postgresql://postgres:postgres@localhost:5432/testdb"
# Additional PostgreSQL environment variables to prevent fallback to
# system user
os.environ["PGUSER"] = "postgres"
os.environ["PGPASSWORD"] = "postgres"
os.environ["PGHOST"] = "localhost"
os.environ["PGPORT"] = "5432"
os.environ["PGDATABASE"] = "testdb"
# ===== STEP 2: Mock the tracker module BEFORE importing app =====
class MockTracker:
def get_tier(self, api_key):
from app.core.usage_tracker import Tier
return Tier.PRO
def get_remaining_quota(self, api_key, tier):
return 1000
def consume_quota_and_log(self, record, idempotency_key=None):
return (True, None)
def increment_usage_sync(self, record, idempotency_key=None):
return True
def get_or_create_api_key(self, key, tier):
return True
def update_api_key_tier(self, key, tier):
return True
def _insert_audit_log(self, record):
pass
# Replace the tracker at the module level
app.core.usage_tracker.tracker = MockTracker()
# ===== STEP 3: Import app and database modules =====
# Force model registration (prevents "no such table" errors)
# Use the environment variable for the database URL (already set)
TEST_DATABASE_URL = os.getenv(
"TEST_DATABASE_URL",
"postgresql://postgres:postgres@localhost:5432/testdb")
if TEST_DATABASE_URL.startswith("postgresql"):
engine = create_engine(TEST_DATABASE_URL)
else:
engine = create_engine(
TEST_DATABASE_URL, connect_args={
"check_same_thread": False})
TestingSessionLocal = sessionmaker(
autocommit=False,
autoflush=False,
bind=engine)
def override_get_db():
db = TestingSessionLocal()
try:
yield db
finally:
db.close()
fastapi_app.dependency_overrides[get_db] = override_get_db
# Override enforce_quota dependency
async def mock_enforce_quota(request: Request, api_key: str = None):
return {"api_key": "test_key", "tier": Tier.PRO, "tenant_id": "test-tenant", "remaining": 1000}
fastapi_app.dependency_overrides[enforce_quota] = mock_enforce_quota
# Override verify_internal_key: production fails closed when
# ARF_INTERNAL_API_KEY is unset, but tests exercise routes directly without
# the gateway-injected X-Internal-Key header.
async def mock_verify_internal_key():
return None
fastapi_app.dependency_overrides[verify_internal_key] = mock_verify_internal_key
@pytest.fixture(scope="session", autouse=True)
def setup_database():
"""Create tables before any tests run."""
Base.metadata.create_all(bind=engine)
yield
Base.metadata.drop_all(bind=engine)
@pytest.fixture(scope="session")
def client():
with TestClient(fastapi_app) as test_client:
yield test_client
@pytest.fixture(scope="function")
def db_session():
"""Provide a database session for each test.
Schema lifecycle is owned entirely by the session-scoped
`setup_database` fixture. Dropping tables here would blow away the
shared schema for any test that runs afterwards without itself
depending on `db_session` (e.g. tests that build their own bare
TestClient), leaving them with a database that has no tables at all.
"""
session = TestingSessionLocal()
yield session
session.rollback()
session.close()
|