school_compare/backend/tests/test_gias_translation.py

"""API-boundary translation: marts now carry GIAS codes; the DataFrame the
rest of the backend sees must carry today's name strings."""

import numpy as np
import pandas as pd

from backend.data_loader import _missing_column_name, translate_gias_code_columns
from backend.gias_codes import ESTABLISHMENT_STATUS, PHASE_OF_EDUCATION


def _code_for(mapping, name):
    return next(c for c, n in mapping.items() if n == name)


def test_codes_become_todays_names():
    df = pd.DataFrame([{
        "urn": 1,
        "phase_code": float(_code_for(PHASE_OF_EDUCATION, "Primary")),
        "school_type_code": np.nan,
        "status_code": float(_code_for(ESTABLISHMENT_STATUS, "Open, but proposed to close")),
        "religious_character_code": np.nan,
        "admissions_policy_code": np.nan,
    }])
    out = translate_gias_code_columns(df)
    row = out.iloc[0]
    assert row["phase"] == "Primary"
    assert row["status"] == "Open, but proposed to close"
    assert row["school_type"] is None
    assert row["religious_denomination"] is None
    assert row["admissions_policy"] is None


def test_unknown_code_degrades_not_blanks():
    df = pd.DataFrame([{"urn": 1, "phase_code": 9999.0}])
    out = translate_gias_code_columns(df)
    assert out.iloc[0]["phase"] == "Unknown (9999)"


def test_missing_code_columns_are_a_noop():
    """Old-schema DataFrames (tests, pre-pipeline DBs) pass through untouched."""
    df = pd.DataFrame([{"urn": 1, "phase": "Primary", "status": "Open"}])
    out = translate_gias_code_columns(df)
    assert out.iloc[0]["phase"] == "Primary"
    assert out.iloc[0]["status"] == "Open"


def _fake_exc(orig_message):
    """A stand-in for sqlalchemy.exc.ProgrammingError: str(exc) embeds the
    full SQL statement (deliberately containing every column name below, to
    prove the matcher doesn't fall back to it), while .orig carries the real
    DBAPI error message naming only the offending column."""
    exc = Exception(
        "SELECT s.phase_code, s.school_type_code, s.religious_character_code, "
        "s.status_code, s.admissions_policy_code, s.has_sixth_form FROM ... "
        f"[SQL: ...] (Background on this error at: https://...)"
    )
    exc.orig = Exception(orig_message) if orig_message is not None else None
    return exc


def test_missing_column_name_quoted():
    assert _missing_column_name(_fake_exc('column "phase_code" does not exist')) == "phase_code"


def test_missing_column_name_unquoted():
    assert _missing_column_name(_fake_exc("column phase_code does not exist")) == "phase_code"


def test_missing_column_name_table_prefixed():
    assert (
        _missing_column_name(_fake_exc("column s.has_sixth_form does not exist"))
        == "has_sixth_form"
    )


def test_missing_column_name_location_alias_prefixed():
    # dim_location columns are selected via the `l.` alias; Postgres reports a
    # missing qualified column unquoted (e.g. "column l.county does not exist").
    # The matcher must strip any alias, not just `s.`, or the county /
    # parliamentary_constituency fallback never triggers and the whole data
    # load degrades to an empty DataFrame (zero schools) instead of NULLs.
    assert (
        _missing_column_name(_fake_exc("column l.county does not exist"))
        == "county"
    )
    assert (
        _missing_column_name(
            _fake_exc("column l.parliamentary_constituency does not exist")
        )
        == "parliamentary_constituency"
    )


def test_missing_column_name_no_match_returns_none():
    assert _missing_column_name(_fake_exc("relation \"marts.dim_school\" does not exist")) is None


def test_load_school_data_survives_premigration_marts(monkeypatch):
    """Real prod state until the nightly pipeline first rebuilds the mart with
    the GIAS code columns: marts.dim_school still has the old name columns
    (phase, school_type, religious_character, status, admissions_policy)
    instead of the new *_code columns. The first query raises UndefinedColumn
    on s.phase_code; load_school_data_as_dataframe must retry with the
    legacy name-column query rather than swallow the error and return (and
    then have load_school_data cache) an empty DataFrame."""
    import sqlalchemy.exc
    from backend import data_loader

    data_loader._df_cache = None
    data_loader._df_latest_cache = None

    good_df = pd.DataFrame(
        [
            {
                "urn": 1,
                "school_name": "Legacy School",
                "phase": "Primary",
                "school_type": "Academy",
                "status": "Open",
            }
        ]
    )
    calls = []

    def fake_read_sql(query, con):
        calls.append(query)
        if len(calls) == 1:
            raise sqlalchemy.exc.ProgrammingError(
                statement=str(data_loader._MAIN_QUERY),
                params=None,
                orig=Exception(
                    "(psycopg2.errors.UndefinedColumn) column s.phase_code "
                    "does not exist\nLINE 5:         s.phase_code,"
                ),
            )
        return good_df.copy()

    monkeypatch.setattr(data_loader.pd, "read_sql", fake_read_sql)

    try:
        df = data_loader.load_school_data_as_dataframe()
    finally:
        data_loader._df_cache = None
        data_loader._df_latest_cache = None

    assert len(calls) == 2, "must retry with the legacy name-column query variant"
    assert calls[1] is data_loader._MAIN_QUERY_LEGACY_NAMES
    assert not df.empty
    assert df["phase"].iloc[0] == "Primary"
    assert df["status"].iloc[0] == "Open"