252 lines
10 KiB
Python
252 lines
10 KiB
Python
import asyncio
|
||
import argparse
|
||
from datetime import date
|
||
from decimal import Decimal
|
||
|
||
from sqlalchemy import delete, select
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
from sqlalchemy.orm import selectinload
|
||
|
||
from app.db.session import async_session_factory
|
||
from app.models.car import Car, CarMake, CarModel, CarTrim
|
||
from app.models.expense import FuelEntry, ServiceEntry, ServiceType
|
||
from app.models.user import User
|
||
from app.services.catalog_data import CAR_CATALOG, CAR_TRIMS, COMMON_TRIMS, MAKE_COUNTRIES
|
||
|
||
|
||
MOCK_PLATE_PREFIX = "MOCK"
|
||
|
||
MOCK_CARS = [
|
||
("KIA Sportage", "KIA", "Sportage", "GT-Line 1.6T DCT AWD", 2021, "gasoline", 36200, Decimal("2450000")),
|
||
("Toyota Camry", "Toyota", "Camry", "Comfort 2.5 AT", 2020, "gasoline", 58400, Decimal("2850000")),
|
||
("Hyundai Tucson", "Hyundai", "Tucson", "Prestige 2.0 AT AWD", 2022, "gasoline", 27100, Decimal("2750000")),
|
||
("Volkswagen Tiguan", "Volkswagen", "Tiguan", "Status 2.0 TSI DSG 4Motion", 2019, "gasoline", 73400, Decimal("2300000")),
|
||
("BMW X3", "BMW", "X3", "20d xDrive", 2021, "diesel", 48900, Decimal("4350000")),
|
||
("Mercedes GLC", "Mercedes", "GLC", "GLC 300 4MATIC AMG Line", 2020, "gasoline", 52200, Decimal("4500000")),
|
||
("Nissan X-Trail", "Nissan", "X-Trail", "Premium", 2018, "gasoline", 91400, Decimal("1850000")),
|
||
("Skoda Octavia", "Skoda", "Octavia", "Comfort", 2021, "gasoline", 46800, Decimal("2050000")),
|
||
("Tesla Model 3", "Tesla", "Model 3", "Long Range AWD", 2022, "electric", 33800, Decimal("3900000")),
|
||
("Haval Jolion", "Haval", "Jolion", "Premium", 2023, "gasoline", 19600, Decimal("2150000")),
|
||
]
|
||
|
||
|
||
def month_shift(base: date, months_back: int) -> date:
|
||
month_index = base.year * 12 + base.month - 1 - months_back
|
||
year = month_index // 12
|
||
month = month_index % 12 + 1
|
||
return date(year, month, min(base.day, 24))
|
||
|
||
|
||
async def seed_catalog(session: AsyncSession) -> None:
|
||
result = await session.execute(select(CarMake).options(selectinload(CarMake.models)))
|
||
existing = {make.name: make for make in result.scalars()}
|
||
|
||
for make_name, model_names in CAR_CATALOG.items():
|
||
make = existing.get(make_name)
|
||
if make is None:
|
||
make = CarMake(name=make_name, country=MAKE_COUNTRIES.get(make_name))
|
||
session.add(make)
|
||
await session.flush()
|
||
existing_models = set()
|
||
else:
|
||
if not make.country and MAKE_COUNTRIES.get(make_name):
|
||
make.country = MAKE_COUNTRIES[make_name]
|
||
existing_models = {model.name for model in make.models}
|
||
for model_name in model_names:
|
||
if model_name not in existing_models:
|
||
session.add(CarModel(make_id=make.id, name=model_name))
|
||
await session.flush()
|
||
await seed_trims(session)
|
||
|
||
|
||
async def seed_trims(session: AsyncSession) -> None:
|
||
result = await session.execute(
|
||
select(CarMake).options(selectinload(CarMake.models).selectinload(CarModel.trims))
|
||
)
|
||
makes = {make.name: make for make in result.scalars()}
|
||
for make in makes.values():
|
||
for model in make.models:
|
||
trim_rows = CAR_TRIMS.get((make.name, model.name)) or [
|
||
{**item, "body_type": infer_body_type(model.name)} for item in COMMON_TRIMS
|
||
]
|
||
existing = {trim.name for trim in model.trims}
|
||
for row in trim_rows:
|
||
if row["name"] not in existing:
|
||
session.add(CarTrim(model_id=model.id, **row))
|
||
|
||
|
||
def infer_body_type(model_name: str) -> str:
|
||
suv_markers = ("Q", "X", "CX", "CR-V", "RAV", "Tiguan", "Tucson", "Sportage", "Jolion")
|
||
return "SUV" if any(marker in model_name for marker in suv_markers) else "sedan"
|
||
|
||
|
||
async def pick_owner(session: AsyncSession) -> User:
|
||
result = await session.execute(select(User).where(User.telegram_id == 1))
|
||
user = result.scalar_one_or_none()
|
||
if user:
|
||
return user
|
||
|
||
user = User(telegram_id=1, username="demo", first_name="Demo")
|
||
session.add(user)
|
||
await session.flush()
|
||
return user
|
||
|
||
|
||
async def clear_previous_mock(session: AsyncSession) -> None:
|
||
result = await session.execute(
|
||
select(Car.id).where(Car.plate_number.like(f"{MOCK_PLATE_PREFIX}-%"))
|
||
)
|
||
car_ids = list(result.scalars())
|
||
if not car_ids:
|
||
return
|
||
await session.execute(delete(ServiceEntry).where(ServiceEntry.car_id.in_(car_ids)))
|
||
await session.execute(delete(FuelEntry).where(FuelEntry.car_id.in_(car_ids)))
|
||
await session.execute(delete(Car).where(Car.id.in_(car_ids)))
|
||
|
||
|
||
async def seed_mock_usage(session: AsyncSession, owner: User) -> None:
|
||
await clear_previous_mock(session)
|
||
today = date.today()
|
||
|
||
for index, (name, make, model, trim, year, fuel_type, start_odo, price) in enumerate(
|
||
MOCK_CARS,
|
||
start=1,
|
||
):
|
||
car = Car(
|
||
owner_id=owner.id,
|
||
name=name,
|
||
make=make,
|
||
model=model,
|
||
trim=trim,
|
||
year=year,
|
||
plate_number=f"{MOCK_PLATE_PREFIX}-{index:02d}",
|
||
fuel_type=fuel_type,
|
||
purchase_date=date(year, min(index, 12), 10),
|
||
purchase_price=price,
|
||
current_odometer=start_odo,
|
||
)
|
||
session.add(car)
|
||
await session.flush()
|
||
|
||
odometer = start_odo
|
||
monthly_km = 820 + index * 95
|
||
consumption = Decimal("7.2") + Decimal(index % 5) * Decimal("0.45")
|
||
if fuel_type == "electric":
|
||
consumption = Decimal("0")
|
||
|
||
for months_back in range(11, -1, -1):
|
||
base_day = month_shift(today.replace(day=15), months_back)
|
||
odometer += monthly_km
|
||
|
||
if fuel_type == "electric":
|
||
energy_cost = Decimal(110 + index * 8 + months_back % 3 * 15)
|
||
session.add(
|
||
ServiceEntry(
|
||
car_id=car.id,
|
||
entry_date=base_day,
|
||
odometer=odometer,
|
||
service_type=ServiceType.other,
|
||
title="Зарядка и парковка",
|
||
category="charging",
|
||
vendor="EV network",
|
||
total_cost=energy_cost,
|
||
)
|
||
)
|
||
else:
|
||
liters_per_month = Decimal(monthly_km) * consumption / Decimal(100)
|
||
price_per_liter = Decimal("58.50") + Decimal(index % 4) * Decimal("2.10")
|
||
for fill in range(2):
|
||
fill_date = date(base_day.year, base_day.month, 8 if fill == 0 else 22)
|
||
liters = (liters_per_month / 2).quantize(Decimal("0.001"))
|
||
session.add(
|
||
FuelEntry(
|
||
car_id=car.id,
|
||
entry_date=fill_date,
|
||
odometer=odometer - (monthly_km // 2 if fill == 0 else 0),
|
||
liters=liters,
|
||
price_per_liter=price_per_liter,
|
||
total_cost=(liters * price_per_liter).quantize(Decimal("0.01")),
|
||
station=["Shell", "Lukoil", "Gazprom", "Rosneft", "Neste"][index % 5],
|
||
is_full_tank=True,
|
||
)
|
||
)
|
||
|
||
if months_back in {11, 5}:
|
||
session.add(
|
||
ServiceEntry(
|
||
car_id=car.id,
|
||
entry_date=date(base_day.year, base_day.month, 12),
|
||
odometer=odometer,
|
||
service_type=ServiceType.maintenance,
|
||
title="Плановое ТО",
|
||
category="regular",
|
||
vendor="Service center",
|
||
total_cost=Decimal(12500 + index * 950),
|
||
next_due_odometer=odometer + 10000,
|
||
)
|
||
)
|
||
|
||
if months_back in {8, 2}:
|
||
session.add(
|
||
ServiceEntry(
|
||
car_id=car.id,
|
||
entry_date=date(base_day.year, base_day.month, 18),
|
||
odometer=odometer,
|
||
service_type=ServiceType.tire,
|
||
title="Сезонная замена шин",
|
||
category="tires",
|
||
vendor="Tire shop",
|
||
total_cost=Decimal(4200 + index * 180),
|
||
)
|
||
)
|
||
|
||
if months_back == 6:
|
||
session.add(
|
||
ServiceEntry(
|
||
car_id=car.id,
|
||
entry_date=date(base_day.year, base_day.month, 20),
|
||
odometer=odometer,
|
||
service_type=ServiceType.insurance,
|
||
title="ОСАГО / страховка",
|
||
category="insurance",
|
||
vendor="Insurance",
|
||
total_cost=Decimal(18200 + index * 1150),
|
||
)
|
||
)
|
||
|
||
if months_back == index % 10:
|
||
session.add(
|
||
ServiceEntry(
|
||
car_id=car.id,
|
||
entry_date=date(base_day.year, base_day.month, 25),
|
||
odometer=odometer,
|
||
service_type=ServiceType.repair,
|
||
title=["Тормозные колодки", "Диагностика подвески", "Замена АКБ", "Ремонт кондиционера"][index % 4],
|
||
category="repair",
|
||
vendor="Garage",
|
||
total_cost=Decimal(7200 + index * 1350),
|
||
)
|
||
)
|
||
|
||
car.current_odometer = odometer
|
||
|
||
|
||
async def main(catalog_only: bool = False) -> None:
|
||
async with async_session_factory() as session:
|
||
await seed_catalog(session)
|
||
if catalog_only:
|
||
await session.commit()
|
||
print("Seeded vehicle catalog")
|
||
return
|
||
owner = await pick_owner(session)
|
||
await seed_mock_usage(session, owner)
|
||
await session.commit()
|
||
print(f"Seeded catalog and mock usage for owner_id={owner.id}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("--catalog-only", action="store_true")
|
||
args = parser.parse_args()
|
||
asyncio.run(main(catalog_only=args.catalog_only))
|