7dffb76801
- SQLTable: add column_meta(), field_descriptions(), to_ddl(), _resolve_type() - generate_models.py: read DuckDB column comments, emit Field(description=...) - run_pipeline: validate output columns against SQLTable contract (SchemaError) - Pipeline.run: pass output class through to runner - Tag.col(): new namespace for column-level descriptions - bib/meta: collect_column_comments() and apply_column_comments() for Zotero→DuckDB flow - stack.toml + src/conf/: centralised config loader with attribute access and path() - 100% test coverage on all changed files
860 lines
29 KiB
Python
860 lines
29 KiB
Python
"""Tests for aco.pipe.runner._param_to_table, run_pipeline, Expr, and Pipeline."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
from unittest.mock import MagicMock
|
|
|
|
import narwhals as nw
|
|
import polars as pl
|
|
import pytest
|
|
|
|
from aco.express.base import Expr
|
|
from aco.pipe.base import Pipeline
|
|
from aco.pipe.runner import SchemaError, _param_to_table, run_pipeline
|
|
from aco.table.base import SQLTable
|
|
from bib.tag import Tag
|
|
|
|
# ── _param_to_table ────────────────────────────────────────────────────────────
|
|
|
|
|
|
class TestParamToTable:
|
|
"""Unit tests for the parameter-name → table-reference converter."""
|
|
|
|
# ── double-underscore (schema.table) ─────────────────────────────────────
|
|
|
|
def test_double_underscore_basic(self) -> None:
|
|
assert _param_to_table("core__encounter") == "core.encounter"
|
|
|
|
def test_double_underscore_schema_and_table(self) -> None:
|
|
assert (
|
|
_param_to_table("input_layer__medical_claim") == "input_layer.medical_claim"
|
|
)
|
|
|
|
def test_double_underscore_claims_preprocessing(self) -> None:
|
|
assert _param_to_table("claims_preprocessing__member_months") == (
|
|
"claims_preprocessing.member_months"
|
|
)
|
|
|
|
def test_double_underscore_readmissions(self) -> None:
|
|
assert _param_to_table("readmissions__encounter") == "readmissions.encounter"
|
|
|
|
def test_double_underscore_splits_on_first_only(self) -> None:
|
|
# schema__table__extra — split only on the FIRST __
|
|
result = _param_to_table("input_layer__input_layer__appointment")
|
|
assert result == "input_layer.input_layer__appointment"
|
|
|
|
def test_double_underscore_pfs_rvu(self) -> None:
|
|
assert _param_to_table("pfs__rvu") == "pfs.rvu"
|
|
|
|
def test_double_underscore_hcc_suspecting(self) -> None:
|
|
assert _param_to_table("hcc_suspecting__int_patient") == (
|
|
"hcc_suspecting.int_patient"
|
|
)
|
|
|
|
# ── triple-underscore (schema._private) ──────────────────────────────────
|
|
|
|
def test_triple_underscore_basic(self) -> None:
|
|
assert _param_to_table("readmissions___int_encounter") == (
|
|
"readmissions._int_encounter"
|
|
)
|
|
|
|
def test_triple_underscore_core(self) -> None:
|
|
assert _param_to_table("core___stg_clinical_encounter") == (
|
|
"core._stg_clinical_encounter"
|
|
)
|
|
|
|
def test_triple_underscore_claims_preprocessing(self) -> None:
|
|
assert _param_to_table("claims_preprocessing___int_medical_claim") == (
|
|
"claims_preprocessing._int_medical_claim"
|
|
)
|
|
|
|
def test_triple_underscore_hcc_suspecting_private(self) -> None:
|
|
assert _param_to_table("hcc_suspecting___int_patient") == (
|
|
"hcc_suspecting._int_patient"
|
|
)
|
|
|
|
# ── bare parameter (no conversion) ───────────────────────────────────────
|
|
|
|
def test_bare_single_word(self) -> None:
|
|
assert _param_to_table("df") == "df"
|
|
|
|
def test_bare_rvu(self) -> None:
|
|
assert _param_to_table("rvu") == "rvu"
|
|
|
|
def test_bare_gpci(self) -> None:
|
|
assert _param_to_table("gpci") == "gpci"
|
|
|
|
def test_bare_claims(self) -> None:
|
|
assert _param_to_table("claims") == "claims"
|
|
|
|
def test_bare_no_underscores(self) -> None:
|
|
assert _param_to_table("labor") == "labor"
|
|
|
|
# ── edge cases ────────────────────────────────────────────────────────────
|
|
|
|
def test_empty_string(self) -> None:
|
|
# Should not raise — return as-is or empty
|
|
result = _param_to_table("")
|
|
assert isinstance(result, str)
|
|
|
|
def test_single_underscore_unchanged(self) -> None:
|
|
# A single underscore in a name should NOT trigger conversion
|
|
result = _param_to_table("some_param")
|
|
assert result == "some_param"
|
|
|
|
def test_output_contains_dot_for_double_underscore(self) -> None:
|
|
result = _param_to_table("schema__table")
|
|
assert "." in result
|
|
|
|
def test_output_contains_dot_for_triple_underscore(self) -> None:
|
|
result = _param_to_table("schema___private")
|
|
assert "." in result
|
|
|
|
def test_triple_takes_priority_over_double(self) -> None:
|
|
"""___ must be detected before __ — both are present in readmissions___int."""
|
|
result = _param_to_table("readmissions___int_encounter")
|
|
# Must be readmissions._int_encounter, NOT readmissions._int_encounter split wrong
|
|
assert result == "readmissions._int_encounter"
|
|
assert result.startswith("readmissions.")
|
|
assert "_int_encounter" in result
|
|
|
|
def test_roundtrip_symmetry(self) -> None:
|
|
"""Different param encodings must resolve to different table refs."""
|
|
double = _param_to_table("core__encounter")
|
|
triple = _param_to_table("core___encounter")
|
|
assert double != triple
|
|
assert double == "core.encounter"
|
|
assert triple == "core._encounter"
|
|
|
|
|
|
# ── run_pipeline ───────────────────────────────────────────────────────────────
|
|
|
|
|
|
class TestRunPipeline:
|
|
"""Unit tests for run_pipeline — the sequential expression executor."""
|
|
|
|
# ── simple cases ──────────────────────────────────────────────────────────
|
|
|
|
def test_single_expression_loaded_from_load(self) -> None:
|
|
"""A single expression whose input comes from load()."""
|
|
source_df = pl.DataFrame({"x": [1, 2, 3]})
|
|
|
|
@nw.narwhalify
|
|
def identity(df):
|
|
return df
|
|
|
|
load = MagicMock(return_value=source_df)
|
|
cache = run_pipeline([("output.table", identity)], load)
|
|
|
|
assert "output.table" in cache
|
|
assert isinstance(cache["output.table"], pl.DataFrame)
|
|
|
|
def test_pipeline_calls_load_for_unknown_inputs(self) -> None:
|
|
"""Inputs not in cache must be fetched via load()."""
|
|
input_df = pl.DataFrame({"a": [10]})
|
|
load = MagicMock(return_value=input_df)
|
|
|
|
@nw.narwhalify
|
|
def passthrough(df):
|
|
return df
|
|
|
|
run_pipeline([("out.table", passthrough)], load)
|
|
load.assert_called_once()
|
|
|
|
def test_pipeline_uses_cache_before_load(self) -> None:
|
|
"""Outputs from earlier steps must be passed to later steps via cache."""
|
|
step1_df = pl.DataFrame({"val": [42]})
|
|
step2_called_with: list[Any] = []
|
|
|
|
@nw.narwhalify
|
|
def step1(df):
|
|
return df
|
|
|
|
@nw.narwhalify
|
|
def step2(step1__result):
|
|
step2_called_with.append(step1__result)
|
|
return step1__result
|
|
|
|
load = MagicMock(return_value=step1_df)
|
|
run_pipeline(
|
|
[("step1.result", step1), ("step2.output", step2)],
|
|
load,
|
|
)
|
|
# step2 received step1's output from cache, not from load
|
|
assert len(step2_called_with) == 1
|
|
|
|
def test_chained_pipeline_produces_both_outputs(self) -> None:
|
|
"""Both step outputs must appear in the returned cache."""
|
|
base = pl.DataFrame({"n": [1, 2]})
|
|
|
|
@nw.narwhalify
|
|
def double(df):
|
|
return df.with_columns(nw.col("n") * 2)
|
|
|
|
@nw.narwhalify
|
|
def triple(step1__doubled):
|
|
return step1__doubled.with_columns(nw.col("n") * 3)
|
|
|
|
load = MagicMock(return_value=base)
|
|
cache = run_pipeline(
|
|
[("step1.doubled", double), ("step2.tripled", triple)],
|
|
load,
|
|
)
|
|
|
|
assert "step1.doubled" in cache
|
|
assert "step2.tripled" in cache
|
|
|
|
def test_empty_pipeline_returns_empty_cache(self) -> None:
|
|
load = MagicMock()
|
|
cache = run_pipeline([], load)
|
|
assert cache == {}
|
|
load.assert_not_called()
|
|
|
|
def test_pipeline_returns_dict(self) -> None:
|
|
df = pl.DataFrame({"x": [1]})
|
|
load = MagicMock(return_value=df)
|
|
|
|
@nw.narwhalify
|
|
def fn(df):
|
|
return df
|
|
|
|
result = run_pipeline([("out.table", fn)], load)
|
|
assert isinstance(result, dict)
|
|
|
|
def test_pipeline_output_name_is_exact(self) -> None:
|
|
"""Cache key must be exactly the output_name string supplied."""
|
|
df = pl.DataFrame({"x": [1]})
|
|
load = MagicMock(return_value=df)
|
|
|
|
@nw.narwhalify
|
|
def fn(df):
|
|
return df
|
|
|
|
cache = run_pipeline([("my_schema.my_table", fn)], load)
|
|
assert "my_schema.my_table" in cache
|
|
|
|
def test_multiple_inputs_resolved(self) -> None:
|
|
"""A function with two inputs should receive both from load/cache."""
|
|
rvu = pl.DataFrame(
|
|
{"hcpcs": ["99213"], "work_rvu": [0.97], "locality": ["0201"]}
|
|
)
|
|
gpci = pl.DataFrame({"locality": ["0201"], "work_gpci": [1.0]})
|
|
|
|
calls: dict[str, Any] = {}
|
|
|
|
@nw.narwhalify
|
|
def join_fn(rvu, gpci):
|
|
calls["rvu"] = rvu
|
|
calls["gpci"] = gpci
|
|
return rvu
|
|
|
|
def load(ref: str):
|
|
return rvu if ref == "rvu" else gpci
|
|
|
|
run_pipeline([("out.joined", join_fn)], load)
|
|
assert "rvu" in calls
|
|
assert "gpci" in calls
|
|
|
|
def test_load_not_called_for_cached_input(self) -> None:
|
|
"""After step1 produces output, step2 should NOT call load for that input."""
|
|
source = pl.DataFrame({"x": [1]})
|
|
load = MagicMock(return_value=source)
|
|
|
|
@nw.narwhalify
|
|
def step1(df):
|
|
return df
|
|
|
|
@nw.narwhalify
|
|
def step2(step1__out):
|
|
return step1__out
|
|
|
|
run_pipeline(
|
|
[("step1.out", step1), ("step2.out", step2)],
|
|
load,
|
|
)
|
|
# load should only be called once (for step1's `df` input)
|
|
load.assert_called_once()
|
|
|
|
def test_pipeline_preserves_execution_order(self) -> None:
|
|
"""Steps must execute in the order given, not dependency-sorted."""
|
|
order: list[str] = []
|
|
|
|
@nw.narwhalify
|
|
def first(df):
|
|
order.append("first")
|
|
return df
|
|
|
|
@nw.narwhalify
|
|
def second(df):
|
|
order.append("second")
|
|
return df
|
|
|
|
load = MagicMock(return_value=pl.DataFrame({"x": [1]}))
|
|
run_pipeline([("a.first", first), ("b.second", second)], load)
|
|
assert order == ["first", "second"]
|
|
|
|
|
|
# ── make_databricks_loader ────────────────────────────────────────────────────
|
|
|
|
|
|
class TestMakeDatabricksLoader:
|
|
def test_returns_callable(self) -> None:
|
|
from aco.pipe.runner import make_databricks_loader
|
|
|
|
spark = MagicMock()
|
|
load = make_databricks_loader(spark, "my_catalog")
|
|
assert callable(load)
|
|
|
|
def test_calls_spark_table_with_catalog_prefix(self) -> None:
|
|
from aco.pipe.runner import make_databricks_loader
|
|
|
|
spark = MagicMock()
|
|
spark.table.return_value = "fake_df"
|
|
load = make_databricks_loader(spark, "cat")
|
|
|
|
result = load("schema.table")
|
|
spark.table.assert_called_once_with("cat.schema.table")
|
|
assert result == "fake_df"
|
|
|
|
|
|
# ── Schema validation ─────────────────────────────────────────────────────────
|
|
|
|
|
|
class TestSchemaValidation:
|
|
"""Tests for output contract validation in run_pipeline."""
|
|
|
|
def test_matching_output_passes(self) -> None:
|
|
class Out(SQLTable):
|
|
__schema__ = "t"
|
|
__tablename__ = "out"
|
|
a: str | None = None
|
|
b: int | None = None
|
|
|
|
@nw.narwhalify
|
|
def fn(t__src):
|
|
return t__src.select("a", "b")
|
|
|
|
src = pl.DataFrame({"a": ["x"], "b": [1], "c": [2]})
|
|
cache = run_pipeline([("t.out", fn, Out)], lambda _: src)
|
|
assert "t.out" in cache
|
|
|
|
def test_missing_column_raises_schema_error(self) -> None:
|
|
class Out(SQLTable):
|
|
__schema__ = "t"
|
|
__tablename__ = "out"
|
|
a: str | None = None
|
|
b: int | None = None
|
|
|
|
@nw.narwhalify
|
|
def fn(t__src):
|
|
return t__src.select("a") # missing b
|
|
|
|
src = pl.DataFrame({"a": ["x"], "b": [1]})
|
|
with pytest.raises(SchemaError, match="missing columns"):
|
|
run_pipeline([("t.out", fn, Out)], lambda _: src)
|
|
|
|
def test_extra_column_raises_schema_error(self) -> None:
|
|
class Out(SQLTable):
|
|
__schema__ = "t"
|
|
__tablename__ = "out"
|
|
a: str | None = None
|
|
|
|
@nw.narwhalify
|
|
def fn(t__src):
|
|
return t__src # has extra column b
|
|
|
|
src = pl.DataFrame({"a": ["x"], "b": [1]})
|
|
with pytest.raises(SchemaError, match="extra columns"):
|
|
run_pipeline([("t.out", fn, Out)], lambda _: src)
|
|
|
|
def test_none_output_skips_validation(self) -> None:
|
|
@nw.narwhalify
|
|
def fn(t__src):
|
|
return t__src
|
|
|
|
src = pl.DataFrame({"any": [1]})
|
|
cache = run_pipeline([("t.out", fn, None)], lambda _: src)
|
|
assert "t.out" in cache
|
|
|
|
def test_two_tuple_skips_validation(self) -> None:
|
|
@nw.narwhalify
|
|
def fn(t__src):
|
|
return t__src
|
|
|
|
src = pl.DataFrame({"any": [1]})
|
|
cache = run_pipeline([("t.out", fn)], lambda _: src)
|
|
assert "t.out" in cache
|
|
|
|
def test_pipeline_run_validates_output(self) -> None:
|
|
"""Pipeline.run passes output classes to runner for validation."""
|
|
|
|
class Out(SQLTable):
|
|
__schema__ = "t"
|
|
__tablename__ = "out"
|
|
a: str | None = None
|
|
|
|
@nw.narwhalify
|
|
def fn(t__src):
|
|
return t__src # has extra column
|
|
|
|
src = pl.DataFrame({"a": ["x"], "b": [1]})
|
|
expr = Expr(name="t.out", fn=fn, output=Out)
|
|
p = Pipeline(exprs=[expr])
|
|
with pytest.raises(SchemaError):
|
|
p.run(lambda _: src)
|
|
|
|
|
|
# ── Expr ───────────────────────────────────────────────────────────────────────
|
|
|
|
|
|
class TestExprConstruction:
|
|
"""Expr must be constructable with valid arguments and reject invalid ones."""
|
|
|
|
@pytest.fixture
|
|
def simple_table(self):
|
|
class MyTable(SQLTable):
|
|
__schema__ = "core"
|
|
__tablename__ = "encounter"
|
|
|
|
return MyTable
|
|
|
|
@pytest.fixture
|
|
def simple_fn(self):
|
|
@nw.narwhalify
|
|
def my_fn(core__encounter):
|
|
"""Return the encounter table."""
|
|
return core__encounter
|
|
|
|
return my_fn
|
|
|
|
# ── valid construction ────────────────────────────────────────────────────
|
|
|
|
def test_minimal_construction(self, simple_fn) -> None:
|
|
expr = Expr(name="core.encounter", fn=simple_fn)
|
|
assert expr.name == "core.encounter"
|
|
assert expr.fn is simple_fn
|
|
|
|
def test_full_construction(self, simple_fn, simple_table) -> None:
|
|
expr = Expr(
|
|
name="core.encounter",
|
|
fn=simple_fn,
|
|
output=simple_table,
|
|
after=["input_layer.encounter"],
|
|
refs=[Tag.module("aco")],
|
|
description="Test expression.",
|
|
)
|
|
assert expr.output is simple_table
|
|
assert expr.after == ["input_layer.encounter"]
|
|
assert len(expr.refs) == 1
|
|
assert expr.description == "Test expression."
|
|
|
|
def test_after_defaults_to_empty_list(self, simple_fn) -> None:
|
|
expr = Expr(name="core.encounter", fn=simple_fn)
|
|
assert expr.after == []
|
|
|
|
def test_refs_defaults_to_empty_list(self, simple_fn) -> None:
|
|
expr = Expr(name="core.encounter", fn=simple_fn)
|
|
assert expr.refs == []
|
|
|
|
def test_output_defaults_to_none(self, simple_fn) -> None:
|
|
expr = Expr(name="core.encounter", fn=simple_fn)
|
|
assert expr.output is None
|
|
|
|
def test_description_defaults_to_empty_string(self, simple_fn) -> None:
|
|
expr = Expr(name="core.encounter", fn=simple_fn)
|
|
assert expr.description == ""
|
|
|
|
# ── name validation ───────────────────────────────────────────────────────
|
|
|
|
def test_name_must_be_qualified(self, simple_fn) -> None:
|
|
"""Unqualified names (no dot) must be rejected."""
|
|
with pytest.raises(Exception): # ValidationError
|
|
Expr(name="encounter", fn=simple_fn)
|
|
|
|
def test_name_with_dot_is_valid(self, simple_fn) -> None:
|
|
expr = Expr(name="readmissions.encounter", fn=simple_fn)
|
|
assert "." in expr.name
|
|
|
|
def test_name_with_private_table(self, simple_fn) -> None:
|
|
expr = Expr(name="core._stg_encounter", fn=simple_fn)
|
|
assert expr.name == "core._stg_encounter"
|
|
|
|
def test_name_with_multi_part_schema(self, simple_fn) -> None:
|
|
# As long as there's at least one dot, it's valid
|
|
expr = Expr(name="input_layer.medical_claim", fn=simple_fn)
|
|
assert expr.name == "input_layer.medical_claim"
|
|
|
|
def test_name_empty_string_rejected(self, simple_fn) -> None:
|
|
with pytest.raises(Exception):
|
|
Expr(name="", fn=simple_fn)
|
|
|
|
def test_name_no_dot_rejected(self, simple_fn) -> None:
|
|
with pytest.raises(Exception):
|
|
Expr(name="just_a_name", fn=simple_fn)
|
|
|
|
|
|
class TestExprInputsProperty:
|
|
"""Expr.inputs derives table refs from the fn's parameter signature."""
|
|
|
|
def test_single_double_underscore_param(self) -> None:
|
|
@nw.narwhalify
|
|
def fn(core__encounter):
|
|
"""fn."""
|
|
return core__encounter
|
|
|
|
expr = Expr(name="out.table", fn=fn)
|
|
assert expr.inputs == ["core.encounter"]
|
|
|
|
def test_multiple_params(self) -> None:
|
|
@nw.narwhalify
|
|
def fn(core__encounter, input_layer__eligibility):
|
|
"""fn."""
|
|
return core__encounter
|
|
|
|
expr = Expr(name="out.table", fn=fn)
|
|
assert "core.encounter" in expr.inputs
|
|
assert "input_layer.eligibility" in expr.inputs
|
|
|
|
def test_bare_param(self) -> None:
|
|
@nw.narwhalify
|
|
def fn(df):
|
|
"""fn."""
|
|
return df
|
|
|
|
expr = Expr(name="out.table", fn=fn)
|
|
assert expr.inputs == ["df"]
|
|
|
|
def test_triple_underscore_param(self) -> None:
|
|
@nw.narwhalify
|
|
def fn(readmissions___int_encounter):
|
|
"""fn."""
|
|
return readmissions___int_encounter
|
|
|
|
expr = Expr(name="out.table", fn=fn)
|
|
assert expr.inputs == ["readmissions._int_encounter"]
|
|
|
|
def test_inputs_preserves_order(self) -> None:
|
|
@nw.narwhalify
|
|
def fn(rvu, gpci, labor):
|
|
"""fn."""
|
|
return rvu
|
|
|
|
expr = Expr(name="out.table", fn=fn)
|
|
assert expr.inputs == ["rvu", "gpci", "labor"]
|
|
|
|
def test_empty_params_yields_empty_inputs(self) -> None:
|
|
@nw.narwhalify
|
|
def fn():
|
|
"""fn."""
|
|
import polars as pl
|
|
|
|
return pl.DataFrame()
|
|
|
|
expr = Expr(name="out.table", fn=fn)
|
|
assert expr.inputs == []
|
|
|
|
|
|
class TestExprDocProperty:
|
|
"""Expr.doc returns description when set, else fn.__doc__."""
|
|
|
|
def test_doc_returns_description_when_set(self) -> None:
|
|
@nw.narwhalify
|
|
def fn(df):
|
|
"""This is the function docstring."""
|
|
return df
|
|
|
|
expr = Expr(name="out.table", fn=fn, description="Custom description.")
|
|
assert expr.doc == "Custom description."
|
|
|
|
def test_doc_falls_back_to_fn_docstring(self) -> None:
|
|
@nw.narwhalify
|
|
def fn(df):
|
|
"""This is the function docstring."""
|
|
return df
|
|
|
|
expr = Expr(name="out.table", fn=fn)
|
|
assert "function docstring" in expr.doc
|
|
|
|
def test_doc_returns_empty_when_no_description_and_no_docstring(self) -> None:
|
|
@nw.narwhalify
|
|
def fn(df):
|
|
return df
|
|
|
|
expr = Expr(name="out.table", fn=fn)
|
|
assert isinstance(expr.doc, str)
|
|
|
|
def test_description_takes_priority_over_docstring(self) -> None:
|
|
@nw.narwhalify
|
|
def fn(df):
|
|
"""Function docstring — should NOT appear."""
|
|
return df
|
|
|
|
expr = Expr(name="out.table", fn=fn, description="Override description.")
|
|
assert expr.doc == "Override description."
|
|
assert "should NOT appear" not in expr.doc
|
|
|
|
|
|
class TestExprQualifiedRefs:
|
|
"""Expr.qualified_refs always includes a table tag auto-generated from name."""
|
|
|
|
def test_qualified_refs_includes_table_tag(self) -> None:
|
|
@nw.narwhalify
|
|
def fn(df):
|
|
"""fn."""
|
|
return df
|
|
|
|
expr = Expr(name="core.encounter", fn=fn)
|
|
labels = [t.label for t in expr.qualified_refs]
|
|
assert "table:core.encounter" in labels
|
|
|
|
def test_qualified_refs_includes_user_refs(self) -> None:
|
|
@nw.narwhalify
|
|
def fn(df):
|
|
"""fn."""
|
|
return df
|
|
|
|
user_tag = Tag.module("aco")
|
|
expr = Expr(name="core.encounter", fn=fn, refs=[user_tag])
|
|
labels = [t.label for t in expr.qualified_refs]
|
|
assert "module:aco" in labels
|
|
|
|
def test_qualified_refs_auto_tag_first(self) -> None:
|
|
@nw.narwhalify
|
|
def fn(df):
|
|
"""fn."""
|
|
return df
|
|
|
|
expr = Expr(name="pfs.rvu", fn=fn)
|
|
# Auto-generated table tag should be first
|
|
assert expr.qualified_refs[0].namespace == "table"
|
|
assert expr.qualified_refs[0].value == "pfs.rvu"
|
|
|
|
def test_qualified_refs_empty_refs_still_has_auto_tag(self) -> None:
|
|
@nw.narwhalify
|
|
def fn(df):
|
|
"""fn."""
|
|
return df
|
|
|
|
expr = Expr(name="aco.eligibility", fn=fn)
|
|
assert len(expr.qualified_refs) == 1
|
|
assert expr.qualified_refs[0].label == "table:aco.eligibility"
|
|
|
|
def test_qualified_refs_multiple_user_tags(self) -> None:
|
|
@nw.narwhalify
|
|
def fn(df):
|
|
"""fn."""
|
|
return df
|
|
|
|
tags = [Tag.module("pfs"), Tag.year(2026), Tag.file("rvu")]
|
|
expr = Expr(name="pfs.rvu", fn=fn, refs=tags)
|
|
labels = [t.label for t in expr.qualified_refs]
|
|
assert "table:pfs.rvu" in labels
|
|
assert "module:pfs" in labels
|
|
assert "year:2026" in labels
|
|
assert "file:rvu" in labels
|
|
|
|
|
|
class TestExprRepr:
|
|
"""Expr.__repr__ must include key fields."""
|
|
|
|
def test_repr_contains_name(self) -> None:
|
|
@nw.narwhalify
|
|
def fn(df):
|
|
"""fn."""
|
|
return df
|
|
|
|
expr = Expr(name="core.encounter", fn=fn)
|
|
assert "core.encounter" in repr(expr)
|
|
|
|
def test_repr_contains_fn_name(self) -> None:
|
|
@nw.narwhalify
|
|
def my_transform(df):
|
|
"""fn."""
|
|
return df
|
|
|
|
expr = Expr(name="core.encounter", fn=my_transform)
|
|
assert "my_transform" in repr(expr)
|
|
|
|
def test_repr_is_string(self) -> None:
|
|
@nw.narwhalify
|
|
def fn(df):
|
|
"""fn."""
|
|
return df
|
|
|
|
expr = Expr(name="core.encounter", fn=fn)
|
|
assert isinstance(repr(expr), str)
|
|
|
|
def test_repr_contains_inputs(self) -> None:
|
|
@nw.narwhalify
|
|
def fn(core__encounter):
|
|
"""fn."""
|
|
return core__encounter
|
|
|
|
expr = Expr(name="out.table", fn=fn)
|
|
r = repr(expr)
|
|
assert "core.encounter" in r
|
|
|
|
def test_repr_contains_output_name_none_when_not_set(self) -> None:
|
|
@nw.narwhalify
|
|
def fn(df):
|
|
"""fn."""
|
|
return df
|
|
|
|
expr = Expr(name="out.table", fn=fn)
|
|
assert "None" in repr(expr)
|
|
|
|
def test_repr_contains_output_class_name_when_set(self) -> None:
|
|
class MyTable(SQLTable):
|
|
__schema__ = "core"
|
|
__tablename__ = "encounter"
|
|
|
|
@nw.narwhalify
|
|
def fn(df):
|
|
"""fn."""
|
|
return df
|
|
|
|
expr = Expr(name="core.encounter", fn=fn, output=MyTable)
|
|
assert "MyTable" in repr(expr)
|
|
|
|
|
|
# ── Pipeline ───────────────────────────────────────────────────────────────────
|
|
|
|
|
|
class TestPipeline:
|
|
"""Tests for Pipeline — ordered sequence of Expr objects."""
|
|
|
|
@pytest.fixture
|
|
def simple_expr(self):
|
|
@nw.narwhalify
|
|
def fn(df):
|
|
"""Simple passthrough."""
|
|
return df
|
|
|
|
return Expr(name="out.table", fn=fn)
|
|
|
|
@pytest.fixture
|
|
def two_exprs(self):
|
|
@nw.narwhalify
|
|
def fn_a(df):
|
|
"""Step A."""
|
|
return df
|
|
|
|
@nw.narwhalify
|
|
def fn_b(df):
|
|
"""Step B."""
|
|
return df
|
|
|
|
return [
|
|
Expr(name="a.table", fn=fn_a),
|
|
Expr(name="b.table", fn=fn_b),
|
|
]
|
|
|
|
# ── construction ──────────────────────────────────────────────────────────
|
|
|
|
def test_pipeline_with_single_expr(self, simple_expr) -> None:
|
|
p = Pipeline(exprs=[simple_expr])
|
|
assert len(p) == 1
|
|
|
|
def test_pipeline_with_multiple_exprs(self, two_exprs) -> None:
|
|
p = Pipeline(exprs=two_exprs)
|
|
assert len(p) == 2
|
|
|
|
def test_empty_pipeline(self) -> None:
|
|
p = Pipeline(exprs=[])
|
|
assert len(p) == 0
|
|
|
|
# ── names() ───────────────────────────────────────────────────────────────
|
|
|
|
def test_names_returns_list(self, simple_expr) -> None:
|
|
p = Pipeline(exprs=[simple_expr])
|
|
assert isinstance(p.names(), list)
|
|
|
|
def test_names_single_expr(self, simple_expr) -> None:
|
|
p = Pipeline(exprs=[simple_expr])
|
|
assert p.names() == ["out.table"]
|
|
|
|
def test_names_multiple_exprs(self, two_exprs) -> None:
|
|
p = Pipeline(exprs=two_exprs)
|
|
assert p.names() == ["a.table", "b.table"]
|
|
|
|
def test_names_empty_pipeline(self) -> None:
|
|
p = Pipeline(exprs=[])
|
|
assert p.names() == []
|
|
|
|
def test_names_preserves_order(self, two_exprs) -> None:
|
|
p = Pipeline(exprs=two_exprs)
|
|
names = p.names()
|
|
assert names[0] == "a.table"
|
|
assert names[1] == "b.table"
|
|
|
|
# ── __len__() ─────────────────────────────────────────────────────────────
|
|
|
|
def test_len_zero(self) -> None:
|
|
assert len(Pipeline(exprs=[])) == 0
|
|
|
|
def test_len_one(self, simple_expr) -> None:
|
|
assert len(Pipeline(exprs=[simple_expr])) == 1
|
|
|
|
def test_len_two(self, two_exprs) -> None:
|
|
assert len(Pipeline(exprs=two_exprs)) == 2
|
|
|
|
# ── run() ─────────────────────────────────────────────────────────────────
|
|
|
|
def test_run_returns_dict(self, simple_expr) -> None:
|
|
df = pl.DataFrame({"x": [1]})
|
|
load = MagicMock(return_value=df)
|
|
p = Pipeline(exprs=[simple_expr])
|
|
result = p.run(load)
|
|
assert isinstance(result, dict)
|
|
|
|
def test_run_output_in_cache(self, simple_expr) -> None:
|
|
df = pl.DataFrame({"x": [1]})
|
|
load = MagicMock(return_value=df)
|
|
p = Pipeline(exprs=[simple_expr])
|
|
result = p.run(load)
|
|
assert "out.table" in result
|
|
|
|
def test_run_empty_pipeline_returns_empty_dict(self) -> None:
|
|
load = MagicMock()
|
|
p = Pipeline(exprs=[])
|
|
result = p.run(load)
|
|
assert result == {}
|
|
|
|
def test_run_chained_pipeline(self) -> None:
|
|
"""Step 2 should receive step 1's output from cache."""
|
|
source = pl.DataFrame({"n": [1, 2, 3]})
|
|
received: list[Any] = []
|
|
|
|
@nw.narwhalify
|
|
def step1(df):
|
|
"""Step 1."""
|
|
return df.with_columns(nw.col("n") + 10)
|
|
|
|
@nw.narwhalify
|
|
def step2(step1__out):
|
|
"""Step 2 consumes step1's output."""
|
|
received.append(step1__out)
|
|
return step1__out
|
|
|
|
exprs = [
|
|
Expr(name="step1.out", fn=step1, after=[]),
|
|
Expr(name="step2.final", fn=step2, after=["step1.out"]),
|
|
]
|
|
p = Pipeline(exprs=exprs)
|
|
load = MagicMock(return_value=source)
|
|
result = p.run(load)
|
|
|
|
assert "step1.out" in result
|
|
assert "step2.final" in result
|
|
# step2 received step1's modified data (n + 10), not the source
|
|
assert len(received) == 1
|
|
|
|
def test_run_all_names_present_in_cache(self, two_exprs) -> None:
|
|
df = pl.DataFrame({"x": [1]})
|
|
load = MagicMock(return_value=df)
|
|
p = Pipeline(exprs=two_exprs)
|
|
result = p.run(load)
|
|
for name in p.names():
|
|
assert name in result
|