504 lines
16 KiB
Python
504 lines
16 KiB
Python
"""Tests for rex.store — persistent state tracking for extraction batches."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pyarrow as pa
|
|
import pytest
|
|
|
|
from rex.store import (
|
|
FILES_SCHEMA,
|
|
PULLS_SCHEMA,
|
|
Store,
|
|
_current_state,
|
|
_generate_key,
|
|
_now,
|
|
)
|
|
|
|
# ── helpers ──────────────────────────────────────────────────────────────────
|
|
|
|
|
|
class TestHelpers:
|
|
def test_now_returns_utc(self) -> None:
|
|
ts = _now()
|
|
assert ts.tzinfo is not None
|
|
|
|
def test_generate_key_length(self) -> None:
|
|
key = _generate_key()
|
|
assert len(key) == 8
|
|
|
|
def test_generate_key_alphanumeric(self) -> None:
|
|
key = _generate_key()
|
|
assert key.isalnum()
|
|
|
|
def test_generate_key_unique(self) -> None:
|
|
keys = {_generate_key() for _ in range(100)}
|
|
# Should be at least 90 unique out of 100
|
|
assert len(keys) > 90
|
|
|
|
|
|
class TestCurrentState:
|
|
def test_empty_table(self) -> None:
|
|
table = PULLS_SCHEMA.empty_table()
|
|
result = _current_state(table, "key")
|
|
assert result.num_rows == 0
|
|
|
|
def test_keeps_last_per_key(self) -> None:
|
|
from datetime import datetime, timezone
|
|
|
|
ts1 = datetime(2024, 1, 1, tzinfo=timezone.utc)
|
|
ts2 = datetime(2024, 1, 2, tzinfo=timezone.utc)
|
|
|
|
table = pa.Table.from_pylist(
|
|
[
|
|
{
|
|
"ts": ts1,
|
|
"key": "K1",
|
|
"format": "fmt",
|
|
"target": "",
|
|
"file_count": 1,
|
|
"rows_total": 0,
|
|
"status": "pending",
|
|
"error": "",
|
|
},
|
|
{
|
|
"ts": ts2,
|
|
"key": "K1",
|
|
"format": "fmt",
|
|
"target": "",
|
|
"file_count": 1,
|
|
"rows_total": 100,
|
|
"status": "complete",
|
|
"error": "",
|
|
},
|
|
],
|
|
schema=PULLS_SCHEMA,
|
|
)
|
|
result = _current_state(table, "key")
|
|
assert result.num_rows == 1
|
|
assert result.column("status")[0].as_py() == "complete"
|
|
|
|
def test_multiple_keys(self) -> None:
|
|
from datetime import datetime, timezone
|
|
|
|
ts = datetime(2024, 1, 1, tzinfo=timezone.utc)
|
|
|
|
table = pa.Table.from_pylist(
|
|
[
|
|
{
|
|
"ts": ts,
|
|
"key": "K1",
|
|
"format": "a",
|
|
"target": "",
|
|
"file_count": 1,
|
|
"rows_total": 0,
|
|
"status": "complete",
|
|
"error": "",
|
|
},
|
|
{
|
|
"ts": ts,
|
|
"key": "K2",
|
|
"format": "b",
|
|
"target": "",
|
|
"file_count": 2,
|
|
"rows_total": 0,
|
|
"status": "pending",
|
|
"error": "",
|
|
},
|
|
],
|
|
schema=PULLS_SCHEMA,
|
|
)
|
|
result = _current_state(table, "key")
|
|
assert result.num_rows == 2
|
|
|
|
|
|
# ── Store state I/O ──────────────────────────────────────────────────────────
|
|
|
|
|
|
class TestStoreState:
|
|
def test_init(self, tmp_path) -> None:
|
|
store = Store(path=str(tmp_path / "rex"))
|
|
assert store._root.endswith("rex")
|
|
|
|
def test_append_and_get_pull(self, tmp_path) -> None:
|
|
store = Store(path=str(tmp_path / "rex"))
|
|
store._append_pull(
|
|
key="TEST01",
|
|
format="sas",
|
|
target="",
|
|
file_count=1,
|
|
status="pending",
|
|
)
|
|
pull = store.get_pull("TEST01")
|
|
assert pull["key"] == "TEST01"
|
|
assert pull["status"] == "pending"
|
|
|
|
def test_get_pull_missing_raises(self, tmp_path) -> None:
|
|
store = Store(path=str(tmp_path / "rex"))
|
|
with pytest.raises(KeyError, match="Pull not found"):
|
|
store.get_pull("MISSING")
|
|
|
|
def test_list_pulls_empty(self, tmp_path) -> None:
|
|
store = Store(path=str(tmp_path / "rex"))
|
|
pulls = store.list_pulls()
|
|
assert pulls == []
|
|
|
|
def test_list_pulls(self, tmp_path) -> None:
|
|
store = Store(path=str(tmp_path / "rex"))
|
|
store._append_pull(key="A1", format="f", status="complete")
|
|
store._append_pull(key="A2", format="f", status="failed")
|
|
pulls = store.list_pulls()
|
|
assert len(pulls) == 2
|
|
|
|
def test_list_pulls_filtered(self, tmp_path) -> None:
|
|
store = Store(path=str(tmp_path / "rex"))
|
|
store._append_pull(key="A1", format="f", status="complete")
|
|
store._append_pull(key="A2", format="f", status="failed")
|
|
complete = store.list_pulls(status="complete")
|
|
assert len(complete) == 1
|
|
assert complete[0]["key"] == "A1"
|
|
|
|
def test_append_and_get_files(self, tmp_path) -> None:
|
|
store = Store(path=str(tmp_path / "rex"))
|
|
store._append_file(
|
|
pull_key="PK1",
|
|
source_path="/data/jan.lst",
|
|
status="complete",
|
|
row_count=100,
|
|
)
|
|
store._append_file(
|
|
pull_key="PK1",
|
|
source_path="/data/feb.lst",
|
|
status="failed",
|
|
error="parse error",
|
|
)
|
|
files = store.get_files("PK1")
|
|
assert len(files) == 2
|
|
complete = [f for f in files if f["status"] == "complete"]
|
|
assert len(complete) == 1
|
|
|
|
def test_get_files_empty(self, tmp_path) -> None:
|
|
store = Store(path=str(tmp_path / "rex"))
|
|
files = store.get_files("NOPULL")
|
|
assert files == []
|
|
|
|
def test_pull_state_updates(self, tmp_path) -> None:
|
|
"""Appending new state for same key returns latest."""
|
|
store = Store(path=str(tmp_path / "rex"))
|
|
store._append_pull(key="U1", format="f", status="pending")
|
|
store._append_pull(key="U1", format="f", status="extracting")
|
|
store._append_pull(key="U1", format="f", status="complete", rows_total=50)
|
|
pull = store.get_pull("U1")
|
|
assert pull["status"] == "complete"
|
|
assert pull["rows_total"] == 50
|
|
|
|
def test_resume_complete_is_noop(self, tmp_path) -> None:
|
|
store = Store(path=str(tmp_path / "rex"))
|
|
store._append_pull(key="R1", format="f", status="complete", rows_total=100)
|
|
result = store.resume("R1")
|
|
assert result == "R1"
|
|
|
|
def test_resume_pending_raises(self, tmp_path) -> None:
|
|
store = Store(path=str(tmp_path / "rex"))
|
|
store._append_pull(key="R2", format="f", status="pending")
|
|
with pytest.raises(ValueError, match="still pending"):
|
|
store.resume("R2")
|
|
|
|
def test_open_method(self, tmp_path) -> None:
|
|
store = Store(path=str(tmp_path / "rex"))
|
|
test_file = tmp_path / "rex" / "test.txt"
|
|
test_file.parent.mkdir(parents=True, exist_ok=True)
|
|
test_file.write_text("hello")
|
|
with store.open(str(test_file), "rb") as f:
|
|
assert f.read() == b"hello"
|
|
|
|
|
|
# ── _store_output ────────────────────────────────────────────────────────────
|
|
|
|
|
|
class TestStoreOutput:
|
|
def test_polars_write(self, tmp_path) -> None:
|
|
import polars as pl
|
|
|
|
store = Store(path=str(tmp_path / "rex"))
|
|
df = pl.DataFrame({"a": [1, 2, 3]})
|
|
dest, size = store._store_output(df, "PK1", "/data/claims.lst")
|
|
assert "claims.parquet" in dest
|
|
assert size > 0
|
|
|
|
def test_non_polars_fallback(self, tmp_path) -> None:
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
store = Store(path=str(tmp_path / "rex"))
|
|
|
|
# Create a fake non-polars object to trigger the else branch
|
|
fake_df = MagicMock()
|
|
fake_table = pa.table({"a": [1, 2]})
|
|
|
|
# Patch at the module level where it's re-imported
|
|
with (
|
|
patch("rex.store.pa.Table") as mock_table_cls,
|
|
patch("rex.store.pq.write_table"),
|
|
):
|
|
mock_table_cls.from_pandas.return_value = fake_table
|
|
dest, size = store._store_output(fake_df, "PK2", "/data/jan.csv")
|
|
|
|
assert "jan.parquet" in dest
|
|
mock_table_cls.from_pandas.assert_called_once_with(fake_df)
|
|
|
|
|
|
# ── extract_batch + _extract_single ──────────────────────────────────────────
|
|
|
|
|
|
class TestExtractBatch:
|
|
def test_successful_extraction(self, tmp_path) -> None:
|
|
from unittest.mock import patch
|
|
|
|
import narwhals as nw
|
|
import polars as pl
|
|
|
|
store = Store(path=str(tmp_path / "rex"))
|
|
mock_df = nw.from_native(pl.DataFrame({"x": [1, 2]}))
|
|
|
|
with (
|
|
patch("rex.store._generate_key", return_value="BATCH01"),
|
|
patch("rex.pipe.extract", return_value=mock_df),
|
|
):
|
|
key = store.extract_batch(
|
|
paths=["/data/a.lst"],
|
|
format="sas",
|
|
target="input_layer.medical_claim",
|
|
)
|
|
|
|
assert key == "BATCH01"
|
|
pull = store.get_pull(key)
|
|
assert pull["status"] == "complete"
|
|
assert pull["rows_total"] == 2
|
|
|
|
files = store.get_files(key)
|
|
assert len(files) == 1
|
|
assert files[0]["status"] == "complete"
|
|
assert files[0]["row_count"] == 2
|
|
|
|
def test_extraction_with_file_failure(self, tmp_path) -> None:
|
|
from unittest.mock import patch
|
|
|
|
store = Store(path=str(tmp_path / "rex"))
|
|
|
|
with (
|
|
patch("rex.store._generate_key", return_value="BATCH02"),
|
|
patch("rex.pipe.extract", side_effect=ValueError("parse error")),
|
|
):
|
|
key = store.extract_batch(
|
|
paths=["/data/bad.lst"],
|
|
format="sas",
|
|
)
|
|
|
|
pull = store.get_pull(key)
|
|
assert pull["status"] == "partial"
|
|
|
|
files = store.get_files(key)
|
|
assert files[0]["status"] == "failed"
|
|
assert "parse error" in files[0]["error"]
|
|
|
|
def test_batch_abort_on_outer_exception(self, tmp_path) -> None:
|
|
"""When _extract_single itself raises (not caught), outer except fires."""
|
|
from unittest.mock import patch
|
|
|
|
store = Store(path=str(tmp_path / "rex"))
|
|
|
|
# Mock _extract_single to raise directly (bypasses inner try/except)
|
|
with (
|
|
patch("rex.store._generate_key", return_value="BATCH03"),
|
|
patch.object(
|
|
store, "_extract_single", side_effect=RuntimeError("outer boom")
|
|
),
|
|
pytest.raises(RuntimeError, match="outer boom"),
|
|
):
|
|
store.extract_batch(
|
|
paths=["/data/a.lst"],
|
|
format="sas",
|
|
)
|
|
|
|
pull = store.get_pull("BATCH03")
|
|
assert pull["status"] == "failed"
|
|
assert "outer boom" in pull["error"]
|
|
|
|
def test_extraction_with_context(self, tmp_path) -> None:
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import narwhals as nw
|
|
import polars as pl
|
|
|
|
store = Store(path=str(tmp_path / "rex"))
|
|
mock_df = nw.from_native(pl.DataFrame({"x": [1]}))
|
|
ctx = MagicMock()
|
|
|
|
with (
|
|
patch("rex.store._generate_key", return_value="BATCH04"),
|
|
patch("rex.pipe.extract", return_value=mock_df),
|
|
):
|
|
store.extract_batch(
|
|
paths=["/data/a.lst"],
|
|
format="sas",
|
|
target="input_layer.medical_claim",
|
|
context=ctx,
|
|
)
|
|
|
|
ctx.save.assert_called_once()
|
|
|
|
def test_extraction_multiple_files(self, tmp_path) -> None:
|
|
from unittest.mock import patch
|
|
|
|
import narwhals as nw
|
|
import polars as pl
|
|
|
|
store = Store(path=str(tmp_path / "rex"))
|
|
mock_df = nw.from_native(pl.DataFrame({"x": [1]}))
|
|
|
|
with (
|
|
patch("rex.store._generate_key", return_value="BATCH05"),
|
|
patch("rex.pipe.extract", return_value=mock_df),
|
|
):
|
|
key = store.extract_batch(
|
|
paths=["/data/a.lst", "/data/b.lst"],
|
|
format="sas",
|
|
)
|
|
|
|
pull = store.get_pull(key)
|
|
assert pull["rows_total"] == 2
|
|
files = store.get_files(key)
|
|
assert len(files) == 2
|
|
|
|
|
|
# ── resume ───────────────────────────────────────────────────────────────────
|
|
|
|
|
|
class TestResume:
|
|
def test_resume_incomplete_files(self, tmp_path) -> None:
|
|
from unittest.mock import patch
|
|
|
|
import narwhals as nw
|
|
import polars as pl
|
|
|
|
store = Store(path=str(tmp_path / "rex"))
|
|
# Simulate a failed pull with one incomplete file
|
|
store._append_pull(
|
|
key="RES1",
|
|
format="sas",
|
|
target="t",
|
|
file_count=2,
|
|
rows_total=5,
|
|
status="failed",
|
|
)
|
|
store._append_file(
|
|
pull_key="RES1",
|
|
source_path="/data/a.lst",
|
|
status="complete",
|
|
row_count=5,
|
|
)
|
|
store._append_file(
|
|
pull_key="RES1",
|
|
source_path="/data/b.lst",
|
|
status="failed",
|
|
error="timeout",
|
|
)
|
|
|
|
mock_df = nw.from_native(pl.DataFrame({"x": [1, 2, 3]}))
|
|
with patch("rex.pipe.extract", return_value=mock_df):
|
|
key = store.resume("RES1")
|
|
|
|
assert key == "RES1"
|
|
pull = store.get_pull(key)
|
|
assert pull["status"] == "complete"
|
|
|
|
def test_resume_all_files_complete(self, tmp_path) -> None:
|
|
store = Store(path=str(tmp_path / "rex"))
|
|
store._append_pull(
|
|
key="RES2",
|
|
format="sas",
|
|
target="t",
|
|
file_count=1,
|
|
rows_total=10,
|
|
status="failed",
|
|
)
|
|
store._append_file(
|
|
pull_key="RES2",
|
|
source_path="/data/a.lst",
|
|
status="complete",
|
|
row_count=10,
|
|
)
|
|
|
|
key = store.resume("RES2")
|
|
assert key == "RES2"
|
|
pull = store.get_pull(key)
|
|
assert pull["status"] == "complete"
|
|
|
|
def test_resume_with_format_override(self, tmp_path) -> None:
|
|
from unittest.mock import patch
|
|
|
|
import narwhals as nw
|
|
import polars as pl
|
|
|
|
store = Store(path=str(tmp_path / "rex"))
|
|
store._append_pull(
|
|
key="RES3",
|
|
format="old_fmt",
|
|
target="t",
|
|
file_count=1,
|
|
status="extracting",
|
|
)
|
|
store._append_file(
|
|
pull_key="RES3",
|
|
source_path="/data/a.lst",
|
|
status="extracting",
|
|
)
|
|
|
|
mock_df = nw.from_native(pl.DataFrame({"x": [1]}))
|
|
with patch("rex.pipe.extract", return_value=mock_df):
|
|
key = store.resume("RES3", format="new_fmt")
|
|
|
|
assert key == "RES3"
|
|
|
|
def test_resume_partial_with_error(self, tmp_path) -> None:
|
|
from unittest.mock import patch
|
|
|
|
store = Store(path=str(tmp_path / "rex"))
|
|
store._append_pull(
|
|
key="RES4",
|
|
format="sas",
|
|
target="t",
|
|
file_count=1,
|
|
status="failed",
|
|
)
|
|
store._append_file(
|
|
pull_key="RES4",
|
|
source_path="/data/a.lst",
|
|
status="failed",
|
|
)
|
|
|
|
with patch("rex.pipe.extract", side_effect=RuntimeError("still broken")):
|
|
key = store.resume("RES4")
|
|
|
|
pull = store.get_pull(key)
|
|
assert pull["status"] == "partial"
|
|
|
|
|
|
# ── Schema constants ─────────────────────────────────────────────────────────
|
|
|
|
|
|
class TestSchemas:
|
|
def test_pulls_schema_fields(self) -> None:
|
|
names = PULLS_SCHEMA.names
|
|
assert "ts" in names
|
|
assert "key" in names
|
|
assert "status" in names
|
|
assert "format" in names
|
|
|
|
def test_files_schema_fields(self) -> None:
|
|
names = FILES_SCHEMA.names
|
|
assert "ts" in names
|
|
assert "pull_key" in names
|
|
assert "source_path" in names
|
|
assert "status" in names
|
|
assert "row_count" in names
|