File size: 5,278 Bytes
fcf8749 | 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 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 |
import pytest
import asyncio
import os
from typing import AsyncGenerator, Generator
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
from sqlalchemy.pool import StaticPool
from httpx import AsyncClient, ASGITransport
from app.database import Base, get_db
from app.main import app
from app.models import Driver, Route, DriverStatsDaily, FairnessConfig
from app.models.driver import VehicleType, PreferredLanguage
from tests.fixtures.test_data import generate_drivers, generate_routes, generate_allocation_request
from datetime import date, timedelta
import numpy as np
# Use in-memory SQLite for tests by default, unless TEST_DATABASE_URL is set
TEST_DB_URL = os.getenv("TEST_DATABASE_URL", "sqlite+aiosqlite:///:memory:")
@pytest.fixture(scope="session")
async def test_engine():
"""Session-scoped test database engine."""
engine = create_async_engine(
TEST_DB_URL,
connect_args={"check_same_thread": False} if "sqlite" in TEST_DB_URL else {},
poolclass=StaticPool if "sqlite" in TEST_DB_URL else None,
)
# Create tables
async with engine.begin() as conn:
await conn.run_sync(Base.metadata.create_all)
yield engine
# Drop tables
async with engine.begin() as conn:
await conn.run_sync(Base.metadata.drop_all)
await engine.dispose()
@pytest.fixture
async def db_session(test_engine) -> AsyncGenerator[AsyncSession, None]:
"""Function-scoped fresh DB session."""
connection = await test_engine.connect()
transaction = await connection.begin()
session_maker = async_sessionmaker(
bind=connection,
class_=AsyncSession,
expire_on_commit=False,
)
session = session_maker()
yield session
await session.close()
await transaction.rollback()
await connection.close()
@pytest.fixture
async def client(db_session) -> AsyncGenerator[AsyncClient, None]:
"""Test client with override for get_db."""
async def override_get_db():
yield db_session
app.dependency_overrides[get_db] = override_get_db
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as ac:
yield ac
app.dependency_overrides.clear()
@pytest.fixture
async def sample_drivers(db_session):
"""50 drivers: 20% EV, mixed experience/stress."""
driver_data = generate_drivers(count=50, ev_ratio=0.2)
drivers = []
for d_data in driver_data:
driver = Driver(
external_id=d_data["id"],
name=d_data["name"],
vehicle_capacity_kg=d_data["vehicle_capacity_kg"],
preferred_language=PreferredLanguage(d_data["preferred_language"]),
vehicle_type=VehicleType.EV if d_data["is_ev"] else VehicleType.ICE,
battery_range_km=d_data["battery_range_km"],
charging_time_minutes=d_data["charging_time_minutes"],
)
db_session.add(driver)
drivers.append(driver)
await db_session.commit()
return drivers
@pytest.fixture
async def sample_routes(db_session):
"""50 routes: varied difficulty."""
# We create routes in DB but allocation request will likely overwrite/use them
# Actually allocation creates routes from packages.
# This fixture is useful if we want to test other things that need existing routes.
# For allocation endpoint testing, we mainly need the request object.
route_data = generate_routes(count=50)
routes = []
for r_data in route_data:
route = Route(
date=date.today(),
cluster_id=r_data["cluster_id"],
total_weight_kg=r_data["total_weight_kg"],
num_packages=r_data["stops"] * 2,
num_stops=r_data["stops"],
route_difficulty_score=r_data["parking_difficulty"],
estimated_time_minutes=r_data["estimated_time_minutes"],
total_distance_km=r_data["total_distance_km"]
)
db_session.add(route)
routes.append(route)
await db_session.commit()
return routes
@pytest.fixture
async def allocation_request(sample_drivers):
"""Complete allocation request payload."""
# We use sample_drivers to ensure IDs match
drivers_list = []
for d in sample_drivers:
drivers_list.append({
"id": d.external_id,
"name": d.name,
"vehicle_capacity_kg": d.vehicle_capacity_kg,
"preferred_language": d.preferred_language.value,
"is_ev": d.vehicle_type == VehicleType.EV
})
# Generate fresh routes for the request
route_data = generate_routes(count=len(drivers_list))
return generate_allocation_request(drivers_list, route_data)
@pytest.fixture
async def active_config(db_session):
"""Active fairness configuration."""
config = FairnessConfig(
is_active=True,
gini_threshold=0.35,
stddev_threshold=25.0,
max_gap_threshold=25.0,
recovery_mode_enabled=True,
ev_safety_margin_pct=10.0,
ev_charging_penalty_weight=0.3,
recovery_penalty_weight=3.0
)
db_session.add(config)
await db_session.commit()
return config
|