Spaces:
Runtime error
Runtime error
File size: 3,486 Bytes
6d20eab 36d5132 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 | """
pytest configuration and fixtures for ARF API tests.
"""
from app.core.usage_tracker import enforce_quota, Tier
from app.api.deps import get_db
from app.database.base import Base
from app.database.models_intents import IntentDB, TenantDB, BetaStateDB, DecisionAuditLogDB # <-- ensure audit table exists
from app.main import app as fastapi_app
from sqlalchemy.orm import sessionmaker
from sqlalchemy import create_engine
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, api_key=None):
return {"api_key": "test_key", "tier": Tier.PRO, "remaining": 1000}
fastapi_app.dependency_overrides[enforce_quota] = mock_enforce_quota
@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 clean database session for each test."""
Base.metadata.create_all(bind=engine)
session = TestingSessionLocal()
yield session
session.rollback()
session.close()
Base.metadata.drop_all(bind=engine)
|