stack/tests/bcda/test_flatten_stream.py

"""bcda.express.flatten — multi-part resolution and streaming Parquet output."""

from __future__ import annotations

import json
import os
from datetime import date
from pathlib import Path

import fsspec
import pyarrow as pa
import pyarrow.parquet as pq

from bcda.express import flatten
from bcda.table.coverage import Coverage
from bcda.table.eob import EobItem, ExplanationOfBenefit
from bcda.table.patient import Patient

ROOT = Path(__file__).resolve().parents[2]
SAMPLES = ROOT / "dev" / "vendor" / "bcda-site"


class TestResolveNdjson:
    def test_parts_sorted_and_later_dirs_win(self, tmp_path: Path) -> None:
        a, b = tmp_path / "a", tmp_path / "b"
        a.mkdir()
        b.mkdir()
        for name in (
            "ExplanationOfBenefit.10.ndjson",
            "ExplanationOfBenefit.ndjson",
            "ExplanationOfBenefit.2.ndjson",
            "ExplanationOfBenefitEmpty.ndjson",
            "Patient.ndjson",
        ):
            (a / name).write_text("{}\n")
        (b / "Patient.ndjson").write_text("{}\n")
        (b / "Coverage.ndjson").write_text("{}\n")
        found = flatten.resolve_ndjson([a, b])
        assert [p.name for p in found["ExplanationOfBenefit"]] == [
            "ExplanationOfBenefit.ndjson",
            "ExplanationOfBenefit.2.ndjson",
            "ExplanationOfBenefit.10.ndjson",
        ]
        assert found["Patient"] == [b / "Patient.ndjson"]
        assert found["Coverage"] == [b / "Coverage.ndjson"]
        assert flatten.resolve_ndjson(str(a))["Patient"] == [a / "Patient.ndjson"]
        assert flatten.resolve_ndjson(tmp_path / "missing") == {}


class TestArrowSchema:
    def test_types(self) -> None:
        s = flatten.arrow_schema(EobItem)
        assert s.field("clm_id").type == pa.string()
        assert s.field("clm_line_num").type == pa.float64()
        assert s.field("clm_line_from_dt").type == pa.date32()
        assert set(s.names) == set(EobItem.column_names())
        for model in (Patient, Coverage, ExplanationOfBenefit):
            assert flatten.arrow_schema(model).names == model.column_names()

    def test_unknown_annotation_falls_back_to_string(self) -> None:
        from conf.table_base import SQLTable

        class Odd(SQLTable):
            weird: dict | None = None
            plain: int = 0

        s = flatten.arrow_schema(Odd)
        assert s.field("weird").type == pa.string()
        assert s.field("plain").type == pa.int64()


class TestParquetSink:
    def test_writes_batches_and_skips_empty(self, tmp_path: Path) -> None:
        fs = fsspec.filesystem("file")
        path = str(tmp_path / "flat" / "p.parquet")
        sink = flatten._ParquetSink(fs, path, flatten.arrow_schema(Patient))
        sink.write([])
        sink.write([{"id": "1", "bene_dob": date(2000, 1, 1)}])
        sink.write([{"id": "2"}])
        sink.close()
        assert sink.rows == 2
        t = pq.read_table(path)
        assert t.num_rows == 2 and t.schema.field("bene_dob").type == pa.date32()
        assert t.column("id").to_pylist() == ["1", "2"]


class TestFlattenExportStreaming:
    def test_official_samples_in_small_batches(
        self, tmp_path: Path, monkeypatch
    ) -> None:
        """Batch size 7 forces many Arrow batches; counts must not change."""
        import shutil

        monkeypatch.setattr(flatten, "BATCH_ROWS", 7)
        src = tmp_path / "ndjson"
        src.mkdir()
        shutil.copy(SAMPLES / "v3_Patient.ndjson", src / "Patient.ndjson")
        shutil.copy(SAMPLES / "v3_Coverage.ndjson", src / "Coverage.ndjson")
        # split the EOB sample into two parts to exercise multi-part reading
        lines = (SAMPLES / "v3_ExplanationOfBenefit.ndjson").read_text().splitlines()
        (src / "ExplanationOfBenefit.ndjson").write_text(
            "\n".join(lines[:100]) + "\n\n"
        )
        (src / "ExplanationOfBenefit.1.ndjson").write_text(
            "\n".join(lines[100:]) + "\n"
        )
        counts = flatten.flatten_export(src, tmp_path / "store")
        assert counts == {
            "bcda.patient": 94,
            "bcda.explanation_of_benefit": 269,
            "bcda.eob_item": 1232,
            "bcda.coverage": 175,
        }
        eob = pq.read_table(
            tmp_path / "store" / "flat" / "explanation_of_benefit.parquet"
        )
        assert eob.num_rows == 269
        assert set(eob.column("system_type").to_pylist()) == {
            "NationalClaimsHistory",
            "DDPS",
        }
        assert eob.schema.field("clm_from_dt").type == pa.date32()

    def test_iter_ndjson_skips_blank_lines(self, tmp_path: Path) -> None:
        p = tmp_path / "x.ndjson"
        p.write_text('{"a":1}\n\n   \n{"a":2}\n')
        assert [r["a"] for r in flatten._iter_ndjson([p])] == [1, 2]
        assert flatten._read_ndjson(p) == [{"a": 1}, {"a": 2}]

    def test_missing_types_write_nothing(self, tmp_path: Path) -> None:
        src = tmp_path / "ndjson"
        src.mkdir()
        (src / "Patient.ndjson").write_text(
            json.dumps(
                {
                    "id": "P",
                    "meta": {
                        "profile": [
                            "http://hl7.org/fhir/us/carin-bb/StructureDefinition/C4BB-Patient|2.2.0"
                        ]
                    },
                }
            )
            + "\n"
        )
        counts = flatten.flatten_export(src, tmp_path / "store")
        assert counts == {"bcda.patient": 1}
        assert sorted(os.listdir(tmp_path / "store" / "flat")) == ["patient.parquet"]