tests/test_database.py

from backend.database import AirQualityDatabase
import pytest


def test_real_dataset_loads_and_answers_a_city_ranking():
    database = AirQualityDatabase()
    database.initialize()

    stats = database.stats()
    assert stats["ready"] is True
    assert stats["records"] > 100_000
    assert stats["cities"] > 100

    columns, rows, truncated = database.execute(
        """
        SELECT city, ROUND(AVG(pm25), 2) AS avg_pm25
        FROM air_quality
        WHERE year = 2023 AND pm25 IS NOT NULL
        GROUP BY city
        ORDER BY avg_pm25 DESC
        LIMIT 10
        """
    )
    assert columns == ["city", "avg_pm25"]
    assert len(rows) == 10
    assert isinstance(rows[0]["avg_pm25"], float)
    assert truncated is False


def test_city_ranking_rigor_requires_equal_station_weighting():
    biased_sql = """
        WITH station_daily AS (
            SELECT city, station, timestamp, AVG(pm25) AS station_pm25
            FROM air_quality
            GROUP BY city, station, timestamp
        )
        SELECT city, AVG(station_pm25) AS avg_pm25
        FROM station_daily
        GROUP BY city
        ORDER BY avg_pm25 DESC
        LIMIT 10
    """
    with pytest.raises(ValueError, match="Analytical rigor check failed"):
        AirQualityDatabase.validate_analytical_rigor(
            "Rank the 10 cities with the highest average PM2.5.",
            biased_sql,
        )


def test_city_ranking_rigor_accepts_station_period_estimates():
    rigorous_sql = """
        WITH station_estimates AS (
            SELECT
                city,
                station,
                AVG(pm25) AS station_pm25,
                COUNT(DISTINCT timestamp) AS observation_days
            FROM air_quality
            WHERE pm25 IS NOT NULL
            GROUP BY city, station
            HAVING COUNT(DISTINCT timestamp) >= 30
        )
        SELECT
            city,
            AVG(station_pm25) AS avg_pm25,
            COUNT(*) AS station_count,
            MIN(observation_days) AS min_observation_days
        FROM station_estimates
        GROUP BY city
        ORDER BY avg_pm25 DESC
        LIMIT 10
    """
    AirQualityDatabase.validate_analytical_rigor(
        "Rank the 10 cities with the highest average PM2.5.",
        rigorous_sql,
    )