oem_crawler/tests/test_raw_storage.py

"""
Tests for RawStorageManager, external compressed storage, and raw extraction toggle.
"""
from __future__ import annotations

import gzip
import json
import shutil
from datetime import datetime, timezone
from decimal import Decimal
from pathlib import Path

import pytest
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from oem_crawler.config import settings
from oem_crawler.crawlers.audi.parser import parse_audi_vehicle
from oem_crawler.crawlers.bmw.parser import parse_bmw_vehicle
from oem_crawler.db.models import Base, Vehicle, VehicleOption
from oem_crawler.db.pipeline import StockIngestionPipeline
from oem_crawler.db.schemas import (
    OptionSchema,
    VehicleCondition,
    VehicleIngestSchema,
    VehicleStatus,
)
from oem_crawler.storage import RawStorageManager


@pytest.fixture
def temp_raw_dir(tmp_path):
    orig_dir = settings.RAW_STORAGE_DIR
    orig_enabled = settings.ENABLE_RAW_EXTRACTION
    raw_dir = tmp_path / "raw_storage"
    settings.RAW_STORAGE_DIR = str(raw_dir)
    settings.ENABLE_RAW_EXTRACTION = True
    yield raw_dir
    settings.RAW_STORAGE_DIR = orig_dir
    settings.ENABLE_RAW_EXTRACTION = orig_enabled


def test_raw_storage_manager_save_and_load(temp_raw_dir):
    payload = {"brand": "BMW", "specs": {"engine": "v8", "power": 390}}
    timestamp = datetime(2026, 11, 7, 12, 0, 0, tzinfo=timezone.utc)

    # 1. Save
    rel_path = RawStorageManager.save_payload(
        brand="BMW",
        oem_vehicle_id="VIN-123456",
        payload=payload,
        timestamp=timestamp,
    )

    assert rel_path == "bmw/2026-11-07/VIN-123456.json.gz"
    full_path = temp_raw_dir / rel_path
    assert full_path.exists()

    # 2. Verify compression
    with gzip.open(full_path, "rt", encoding="utf-8") as f:
        loaded_raw = json.load(f)
    assert loaded_raw == payload

    # 3. Load via Manager
    loaded = RawStorageManager.load_payload(rel_path)
    assert loaded == payload

    # 4. Delete
    deleted = RawStorageManager.delete_payload(rel_path)
    assert deleted is True
    assert not full_path.exists()


def test_raw_storage_manager_disabled(temp_raw_dir):
    settings.ENABLE_RAW_EXTRACTION = False
    payload = {"key": "val"}

    rel_path = RawStorageManager.save_payload(
        brand="AUDI",
        oem_vehicle_id="AUDI-001",
        payload=payload,
    )
    assert rel_path is None
    assert not (temp_raw_dir / "audi").exists()


def test_parser_include_raw_toggle():
    # BMW parser with include_raw=False
    raw_bmw = {
        "documentId": "TEST-BMW-RAW",
        "vehicleSpecification": {
            "modelAndOption": {
                "model": {"modelName": "BMW 330i"},
                "modelRange": {"description": {"de_DE": "3er"}},
            }
        },
        "pricing": {"price": {"value": 55000}},
        "extra_huge_catalog": {"items": [1, 2, 3] * 1000},
    }

    veh_with_raw = parse_bmw_vehicle(raw_bmw, include_raw=True)
    assert veh_with_raw.raw_payload == raw_bmw

    veh_no_raw = parse_bmw_vehicle(raw_bmw, include_raw=False)
    assert veh_no_raw.raw_payload == {}
    assert veh_no_raw.model_name == "BMW 330i"
    assert veh_no_raw.series == "3er"


def test_pipeline_ingestion_with_raw_storage(temp_raw_dir, tmp_path):
    db_file = tmp_path / "test_stock.db"
    db_url = f"sqlite:///{db_file}"
    engine = create_engine(db_url)
    Base.metadata.create_all(engine)
    session_factory = sessionmaker(bind=engine)

    pipeline = StockIngestionPipeline(session_factory=session_factory)

    test_vehicle = VehicleIngestSchema(
        brand="BMW",
        oem_vehicle_id="BMW-INGEST-01",
        series="3er",
        model_name="320d",
        condition=VehicleCondition.NEW,
        fuel_type="diesel",
        mileage_km=0,
        price_gross=Decimal("45000.00"),
        status=VehicleStatus.IN_STOCK,
        options=[
            OptionSchema(option_code="S01", name="Sport Line", raw_data={"pr": "S01"}),
        ],
        raw_payload={"full_spec": "all details here", "code": "BMW-INGEST-01"},
    )

    def stream():
        yield test_vehicle

    run = pipeline.run_crawler_pipeline(
        brand="BMW",
        crawler_name="test-crawler",
        vehicle_stream=stream(),
        market_country="DE",
    )
    assert run.items_inserted == 1

    with session_factory() as session:
        veh = session.query(Vehicle).filter_by(oem_vehicle_id="BMW-INGEST-01").first()
        assert veh is not None
        # In-database raw_payload is kept empty to save disk space
        assert veh.raw_payload == {}
        # Relative path points to the compressed file
        assert veh.raw_payload_path is not None
        assert veh.raw_payload_path.startswith("bmw/")
        assert veh.raw_payload_path.endswith(".json.gz")

        # Options raw_data in database is kept empty
        opt = session.query(VehicleOption).filter_by(vehicle_id=veh.id).first()
        assert opt.name == "Sport Line"
        assert opt.raw_data == {}

        # Loading via helper method reads from disk seamlessly
        recovered_payload = veh.get_raw_payload()
        assert recovered_payload == {"full_spec": "all details here", "code": "BMW-INGEST-01"}


def test_restore_database_from_raw(temp_raw_dir, tmp_path):
    from oem_crawler.tools.restore_from_raw import restore_database

    # Create dummy raw payload files in temp_raw_dir
    payload = {
        "documentId": "BMW-RESTORE-99",
        "vehicleSpecification": {
            "modelAndOption": {
                "model": {"modelName": "BMW 520d"},
                "modelRange": {"description": {"de_DE": "5er"}},
            }
        },
        "pricing": {"price": {"value": 62000}},
    }
    RawStorageManager.save_payload(
        brand="bmw",
        oem_vehicle_id="BMW-RESTORE-99",
        payload=payload,
    )

    db_file = tmp_path / "restored.db"
    db_url = f"sqlite:///{db_file}"

    # Restore into fresh DB
    restore_database(
        raw_dir=temp_raw_dir,
        db_url=db_url,
        target_brands=["bmw"],
    )

    engine = create_engine(db_url)
    session_factory = sessionmaker(bind=engine)
    with session_factory() as session:
        veh = session.query(Vehicle).filter_by(oem_vehicle_id="BMW-RESTORE-99").first()
        assert veh is not None
        assert veh.model_name == "BMW 520d"
        assert veh.price_gross == Decimal("62000")
        assert veh.raw_payload_path.startswith("bmw/")
        assert veh.raw_payload == {}