Files
stack/tests/aco/test_express_main.py

656 lines
22 KiB
Python

"""Tests for aco.express.main — main schema transformation functions.
Tests verify filtering logic (internal flag, status filters), union
behavior, column renames, and group-by aggregations.
"""
from __future__ import annotations
from datetime import datetime
import polars as pl
import pytest
from aco.express.main import (
alerts_anomaly_detection,
alerts_dbt_models,
alerts_dbt_source_freshness,
alerts_dbt_tests,
alerts_schema_changes,
anomaly_threshold_sensitivity,
dbt_artifacts_hashes,
duckdb_columns,
duckdb_constraints,
duckdb_databases,
duckdb_indexes,
duckdb_logs,
duckdb_schemas,
duckdb_tables,
duckdb_types,
duckdb_views,
hcc_suspecting__list_all,
job_run_results,
metrics_anomaly_score,
model_run_results,
monitors_runs,
pragma_database_list,
seed_run_results,
snapshot_run_results,
sqlite_master,
sqlite_schema,
sqlite_temp_master,
sqlite_temp_schema,
)
# ── helpers ──────────────────────────────────────────────────────
def _cols(result: pl.DataFrame, *expected: str) -> None:
missing = [c for c in expected if c not in result.columns]
assert not missing, f"Missing columns: {missing}"
# ── alerts_dbt_models ────────────────────────────────────────────
def _make_run_results(**overrides) -> pl.DataFrame:
"""Build a minimal run-results frame with all 23 common cols."""
base = {
"model_execution_id": ["exec1"],
"unique_id": ["model.pkg.my_model"],
"invocation_id": ["inv1"],
"name": ["my_model"],
"generated_at": [datetime(2024, 6, 1, 12, 0)],
"status": ["success"],
"full_refresh": [False],
"message": ["OK"],
"execution_time": [1.5],
"execute_started_at": [datetime(2024, 6, 1, 12, 0)],
"execute_completed_at": [datetime(2024, 6, 1, 12, 1)],
"compile_started_at": [datetime(2024, 6, 1, 11, 59)],
"compile_completed_at": [datetime(2024, 6, 1, 12, 0)],
"compiled_code": ["SELECT 1"],
"database_name": ["analytics"],
"schema_name": ["public"],
"materialization": ["table"],
"tags": ["daily"],
"package_name": ["pkg"],
"path": ["models/my_model.sql"],
"original_path": ["models/my_model.sql"],
"owner": ["team_a"],
"alias": ["my_model"],
}
base.update(overrides)
return pl.DataFrame(base)
class TestAlertsDbtModels:
"""alerts_dbt_models unions two frames and keeps
non-success/non-skipped rows."""
def test_returns_dataframe(self) -> None:
m = _make_run_results()
s = _make_run_results()
result = alerts_dbt_models(m, s)
assert isinstance(result, pl.DataFrame)
def test_filters_out_success(self) -> None:
m = _make_run_results(status=["success"])
s = _make_run_results(status=["success"])
result = alerts_dbt_models(m, s)
assert len(result) == 0
def test_filters_out_skipped(self) -> None:
m = _make_run_results(status=["skipped"])
s = _make_run_results(status=["Skipped"])
result = alerts_dbt_models(m, s)
assert len(result) == 0
def test_keeps_error_status(self) -> None:
m = _make_run_results(status=["error"])
s = _make_run_results(status=["success"])
result = alerts_dbt_models(m, s)
assert len(result) == 1
assert result["status"][0] == "error"
def test_unions_both_sources(self) -> None:
m = _make_run_results(
model_execution_id=["e1"],
status=["error"],
)
s = _make_run_results(
model_execution_id=["e2"],
status=["fail"],
)
result = alerts_dbt_models(m, s)
assert len(result) == 2
def test_renames_columns(self) -> None:
m = _make_run_results(status=["error"])
s = _make_run_results(status=["success"])
result = alerts_dbt_models(m, s)
_cols(
result,
"alert_id",
"detected_at",
"owners",
)
assert "model_execution_id" not in result.columns
assert "generated_at" not in result.columns
assert "owner" not in result.columns
def test_alert_id_from_model_execution_id(self) -> None:
m = _make_run_results(
model_execution_id=["abc123"],
status=["error"],
)
s = _make_run_results(status=["success"])
result = alerts_dbt_models(m, s)
assert result["alert_id"][0] == "abc123"
def test_case_insensitive_filter(self) -> None:
m = _make_run_results(status=["SUCCESS"])
s = _make_run_results(status=["Success"])
result = alerts_dbt_models(m, s)
assert len(result) == 0
def test_empty_inputs(self) -> None:
m = _make_run_results().head(0)
s = _make_run_results().head(0)
result = alerts_dbt_models(m, s)
assert len(result) == 0
# ── alerts_dbt_source_freshness ──────────────────────────────────
class TestAlertsDbtSourceFreshness:
"""alerts_dbt_source_freshness joins freshness with
sources and filters non-pass status."""
@pytest.fixture
def freshness_df(self) -> pl.DataFrame:
return pl.DataFrame(
{
"source_freshness_execution_id": [
"sf1",
"sf2",
],
"unique_id": ["src.a", "src.b"],
"max_loaded_at": [
datetime(2024, 6, 1),
datetime(2024, 6, 1),
],
"snapshotted_at": [
datetime(2024, 6, 2),
datetime(2024, 6, 2),
],
"generated_at": [
datetime(2024, 6, 2),
datetime(2024, 6, 2),
],
"max_loaded_at_time_ago_in_s": [
86400.0,
86400.0,
],
"status": ["warn", "pass"],
"error": [None, None],
"warn_after": ["24h", "24h"],
"error_after": ["48h", "48h"],
"filter": [None, None],
}
)
@pytest.fixture
def sources_df(self) -> pl.DataFrame:
return pl.DataFrame(
{
"unique_id": ["src.a", "src.b"],
"database_name": ["db", "db"],
"schema_name": ["raw", "raw"],
"source_name": ["events", "users"],
"identifier": ["events", "users"],
"tags": ["daily", "daily"],
"meta": ["{}", "{}"],
"owner": ["team", "team"],
"package_name": ["pkg", "pkg"],
"path": ["sources.yml", "sources.yml"],
}
)
def test_returns_dataframe(self, freshness_df, sources_df) -> None:
result = alerts_dbt_source_freshness(freshness_df, sources_df)
assert isinstance(result, pl.DataFrame)
def test_filters_out_pass(self, freshness_df, sources_df) -> None:
result = alerts_dbt_source_freshness(freshness_df, sources_df)
assert len(result) == 1
assert result["unique_id"][0] == "src.a"
def test_renames_alert_id(self, freshness_df, sources_df) -> None:
result = alerts_dbt_source_freshness(freshness_df, sources_df)
assert "alert_id" in result.columns
assert result["alert_id"][0] == "sf1"
def test_renames_detected_at(self, freshness_df, sources_df) -> None:
result = alerts_dbt_source_freshness(freshness_df, sources_df)
assert "detected_at" in result.columns
def test_no_match_returns_empty(self) -> None:
fr = pl.DataFrame(
{
"source_freshness_execution_id": ["x"],
"unique_id": ["no_match"],
"max_loaded_at": [datetime(2024, 1, 1)],
"snapshotted_at": [datetime(2024, 1, 1)],
"generated_at": [datetime(2024, 1, 1)],
"max_loaded_at_time_ago_in_s": [0.0],
"status": ["warn"],
"error": [None],
"warn_after": ["1h"],
"error_after": ["2h"],
"filter": [None],
}
)
src = pl.DataFrame(
{
"unique_id": ["other"],
"database_name": ["db"],
"schema_name": ["raw"],
"source_name": ["tbl"],
"identifier": ["tbl"],
"tags": [""],
"meta": ["{}"],
"owner": ["t"],
"package_name": ["p"],
"path": ["s.yml"],
}
)
result = alerts_dbt_source_freshness(fr, src)
assert len(result) == 0
# ── dbt_artifacts_hashes ─────────────────────────────────────────
class TestDbtArtifactsHashes:
"""dbt_artifacts_hashes unions (label, metadata_hash) from
8 artifact tables."""
def _make_artifact(self, hashes):
return pl.DataFrame({"metadata_hash": hashes})
def test_unions_all_eight(self) -> None:
arts = [self._make_artifact(["h1"]) for _ in range(8)]
result = dbt_artifacts_hashes(*arts)
assert len(result) == 8
def test_labels_present(self) -> None:
arts = [self._make_artifact(["h"]) for _ in range(8)]
result = dbt_artifacts_hashes(*arts)
labels = sorted(result["artifacts_model"].to_list())
expected = sorted(
[
"dbt_models",
"dbt_tests",
"dbt_sources",
"dbt_snapshots",
"dbt_metrics",
"dbt_exposures",
"dbt_seeds",
"dbt_columns",
]
)
assert labels == expected
def test_output_columns(self) -> None:
arts = [self._make_artifact(["h"]) for _ in range(8)]
result = dbt_artifacts_hashes(*arts)
assert set(result.columns) == {
"artifacts_model",
"metadata_hash",
}
def test_multiple_hashes_per_artifact(self) -> None:
arts = [self._make_artifact(["h1", "h2"])] + [
self._make_artifact(["h"]) for _ in range(7)
]
result = dbt_artifacts_hashes(*arts)
assert len(result) == 9
# ── duckdb internal filter functions ─────────────────────────────
class TestDuckdbInternalFilter:
"""duckdb_columns, databases, schemas, tables, views
all filter where ~internal."""
@pytest.fixture
def df_with_internal(self) -> pl.DataFrame:
return pl.DataFrame(
{
"name": [
"my_table",
"sys_table",
"other",
],
"internal": [False, True, False],
}
)
def test_duckdb_columns_filters(self, df_with_internal) -> None:
result = duckdb_columns(df_with_internal)
assert len(result) == 2
assert "sys_table" not in result["name"].to_list()
def test_duckdb_databases_filters(self, df_with_internal) -> None:
result = duckdb_databases(df_with_internal)
assert len(result) == 2
def test_duckdb_schemas_filters(self, df_with_internal) -> None:
result = duckdb_schemas(df_with_internal)
assert len(result) == 2
def test_duckdb_tables_filters(self, df_with_internal) -> None:
result = duckdb_tables(df_with_internal)
assert len(result) == 2
def test_duckdb_views_filters(self, df_with_internal) -> None:
result = duckdb_views(df_with_internal)
assert len(result) == 2
def test_all_internal_returns_empty(self) -> None:
df = pl.DataFrame(
{
"name": ["a", "b"],
"internal": [True, True],
}
)
for fn in [
duckdb_columns,
duckdb_databases,
duckdb_schemas,
duckdb_tables,
duckdb_views,
]:
assert len(fn(df)) == 0
def test_no_internal_returns_all(self) -> None:
df = pl.DataFrame(
{
"name": ["a", "b"],
"internal": [False, False],
}
)
for fn in [
duckdb_columns,
duckdb_databases,
duckdb_schemas,
duckdb_tables,
duckdb_views,
]:
assert len(fn(df)) == 2
def test_pragma_database_list_filters(self, df_with_internal) -> None:
result = pragma_database_list(df_with_internal)
assert len(result) == 2
names = result["name"].to_list()
assert "sys_table" not in names
# ── hcc_suspecting__list_all ─────────────────────────────────────
class TestHccSuspectingListAll:
"""hcc_suspecting__list_all unions 5 frames on 10
common columns and deduplicates."""
def _make_suspects(self, ids):
n = len(ids)
return pl.DataFrame(
{
"person_id": ids,
"payer": ["Medicare"] * n,
"data_source": ["test"] * n,
"model_version": ["v1"] * n,
"hcc_code": ["HCC19"] * n,
"hcc_description": ["Diabetes"] * n,
"reason": ["history"] * n,
"contributing_factor": ["dx"] * n,
"suspect_date": ["2024-01-01"] * n,
"current_year_billed": [False] * n,
}
)
def test_unions_five_frames(self) -> None:
frames = [
self._make_suspects(["P1"]),
self._make_suspects(["P2"]),
self._make_suspects(["P3"]),
self._make_suspects(["P4"]),
self._make_suspects(["P5"]),
]
result = hcc_suspecting__list_all(*frames)
assert len(result) == 5
def test_deduplicates(self) -> None:
same = self._make_suspects(["P1"])
frames = [same] * 5
result = hcc_suspecting__list_all(*frames)
assert len(result) == 1
def test_output_columns(self) -> None:
frames = [self._make_suspects(["P1"]) for _ in range(5)]
result = hcc_suspecting__list_all(*frames)
expected = {
"person_id",
"payer",
"data_source",
"model_version",
"hcc_code",
"hcc_description",
"reason",
"contributing_factor",
"suspect_date",
"current_year_billed",
}
assert set(result.columns) == expected
def test_partial_duplicates(self) -> None:
f1 = self._make_suspects(["P1", "P2"])
f2 = self._make_suspects(["P2", "P3"])
f3 = self._make_suspects(["P4"])
f4 = self._make_suspects(["P5"])
f5 = self._make_suspects(["P1"])
result = hcc_suspecting__list_all(f1, f2, f3, f4, f5)
assert len(result) == 5
# ── job_run_results ──────────────────────────────────────────────
class TestJobRunResults:
"""job_run_results filters non-null job_id, groups by
job_name/job_id/job_run_id, aggs min/max dates."""
@pytest.fixture
def invocations_df(self) -> pl.DataFrame:
return pl.DataFrame(
{
"job_name": [
"nightly",
"nightly",
"adhoc",
"manual",
],
"job_id": [
"J1",
"J1",
"J2",
None,
],
"job_run_id": [
"R1",
"R1",
"R2",
"R3",
],
"run_started_at": [
datetime(2024, 6, 1, 8, 0),
datetime(2024, 6, 1, 8, 5),
datetime(2024, 6, 1, 9, 0),
datetime(2024, 6, 1, 10, 0),
],
"run_completed_at": [
datetime(2024, 6, 1, 8, 10),
datetime(2024, 6, 1, 8, 20),
datetime(2024, 6, 1, 9, 30),
datetime(2024, 6, 1, 10, 15),
],
}
)
def test_returns_dataframe(self, invocations_df) -> None:
result = job_run_results(invocations_df)
assert isinstance(result, pl.DataFrame)
def test_filters_null_job_id(self, invocations_df) -> None:
result = job_run_results(invocations_df)
assert len(result) == 2
ids = result["id"].to_list()
assert None not in ids
def test_groups_and_aggs(self, invocations_df) -> None:
result = job_run_results(invocations_df)
nightly = result.filter(pl.col("name") == "nightly")
assert len(nightly) == 1
assert nightly["run_started_at"][0] == datetime(2024, 6, 1, 8, 0)
assert nightly["run_completed_at"][0] == datetime(2024, 6, 1, 8, 20)
def test_renames_output_columns(self, invocations_df) -> None:
result = job_run_results(invocations_df)
_cols(
result,
"name",
"id",
"run_id",
"run_started_at",
"run_completed_at",
)
assert "job_name" not in result.columns
assert "job_id" not in result.columns
def test_all_null_job_ids(self) -> None:
df = pl.DataFrame(
{
"job_name": ["a"],
"job_id": [None],
"job_run_id": ["r"],
"run_started_at": [datetime(2024, 1, 1)],
"run_completed_at": [datetime(2024, 1, 1)],
}
)
result = job_run_results(df)
assert len(result) == 0
def test_single_invocation_per_group(self) -> None:
df = pl.DataFrame(
{
"job_name": ["build"],
"job_id": ["J1"],
"job_run_id": ["R1"],
"run_started_at": [datetime(2024, 1, 1, 9, 0)],
"run_completed_at": [datetime(2024, 1, 1, 9, 30)],
}
)
result = job_run_results(df)
assert len(result) == 1
assert result["name"][0] == "build"
assert result["id"][0] == "J1"
assert result["run_id"][0] == "R1"
# ── passthrough / simple functions ────────────────────────────────
class TestPassthroughFunctions:
"""Simple passthrough or single-input functions that return
their input directly or with minimal transformation."""
@pytest.fixture
def simple_df(self) -> pl.DataFrame:
return pl.DataFrame({"x": [1, 2], "y": ["a", "b"]})
def test_alerts_anomaly_detection(self, simple_df) -> None:
result = alerts_anomaly_detection(simple_df)
assert result.shape == simple_df.shape
def test_alerts_dbt_tests(self, simple_df) -> None:
result = alerts_dbt_tests(simple_df)
assert result.shape == simple_df.shape
def test_alerts_schema_changes(self, simple_df) -> None:
result = alerts_schema_changes(simple_df)
assert result.shape == simple_df.shape
def test_anomaly_threshold_sensitivity(self, simple_df) -> None:
result = anomaly_threshold_sensitivity(simple_df)
assert result.shape == simple_df.shape
def test_duckdb_constraints(self, simple_df) -> None:
result = duckdb_constraints(simple_df)
assert result.shape == simple_df.shape
def test_duckdb_indexes(self, simple_df) -> None:
result = duckdb_indexes(simple_df)
assert result.shape == simple_df.shape
def test_duckdb_logs(self, simple_df) -> None:
result = duckdb_logs(simple_df)
assert result.shape == simple_df.shape
def test_duckdb_types(self, simple_df) -> None:
result = duckdb_types(simple_df)
assert result.shape == simple_df.shape
def test_metrics_anomaly_score(self, simple_df) -> None:
result = metrics_anomaly_score(simple_df)
assert result.shape == simple_df.shape
def test_model_run_results(self, simple_df) -> None:
# Takes two args; returns first
df2 = pl.DataFrame({"z": [1]})
result = model_run_results(simple_df, df2)
assert result.shape == simple_df.shape
def test_monitors_runs(self, simple_df) -> None:
result = monitors_runs(simple_df)
assert result.shape == simple_df.shape
def test_seed_run_results(self, simple_df) -> None:
df2 = pl.DataFrame({"z": [1]})
result = seed_run_results(simple_df, df2)
assert result.shape == simple_df.shape
def test_snapshot_run_results(self, simple_df) -> None:
df2 = pl.DataFrame({"z": [1]})
result = snapshot_run_results(simple_df, df2)
assert result.shape == simple_df.shape
def test_sqlite_master(self, simple_df) -> None:
result = sqlite_master(simple_df)
assert result.shape == simple_df.shape
def test_sqlite_schema(self, simple_df) -> None:
result = sqlite_schema(simple_df)
assert result.shape == simple_df.shape
def test_sqlite_temp_master(self, simple_df) -> None:
result = sqlite_temp_master(simple_df)
assert result.shape == simple_df.shape
def test_sqlite_temp_schema(self, simple_df) -> None:
result = sqlite_temp_schema(simple_df)
assert result.shape == simple_df.shape