stack/tests/bcda/test_store.py

"""Tests for bcda.store — persistent storage and state management."""

from __future__ import annotations

from pathlib import Path
from unittest.mock import MagicMock

import pyarrow as pa
import pyarrow.compute  # noqa: F401  — ensure pa.compute is loaded
import pyarrow.parquet as pq
import pytest

from bcda.store import (
    FILES_SCHEMA,
    JOBS_SCHEMA,
    Store,
    _current_state,
    _generate_key,
    _now,
)

# ── Module-level helpers ─────────────────────────────────────


class TestHelpers:
    def test_now_returns_utc(self) -> None:
        ts = _now()
        assert ts.tzinfo is not None
        assert ts.tzname() == "UTC"

    def test_generate_key_length(self) -> None:
        key = _generate_key()
        assert len(key) == 8
        assert key.isalnum()

    def test_generate_key_unique(self) -> None:
        keys = {_generate_key() for _ in range(50)}
        # Extremely unlikely to get <40 unique keys from 50 trials.
        assert len(keys) > 40

    def test_current_state_empty(self) -> None:
        table = JOBS_SCHEMA.empty_table()
        result = _current_state(table, "key")
        assert result.num_rows == 0

    def test_current_state_dedup(self) -> None:
        from datetime import datetime, timezone

        rows = [
            {
                "ts": datetime(2025, 1, 1, tzinfo=timezone.utc),
                "key": "A",
                "job_url": "",
                "endpoint": "",
                "types": "",
                "since": "",
                "status": "pending",
                "error": "",
            },
            {
                "ts": datetime(2025, 1, 2, tzinfo=timezone.utc),
                "key": "A",
                "job_url": "url",
                "endpoint": "",
                "types": "",
                "since": "",
                "status": "complete",
                "error": "",
            },
            {
                "ts": datetime(2025, 1, 1, tzinfo=timezone.utc),
                "key": "B",
                "job_url": "",
                "endpoint": "",
                "types": "",
                "since": "",
                "status": "failed",
                "error": "err",
            },
        ]
        table = pa.Table.from_pylist(rows, schema=JOBS_SCHEMA)
        result = _current_state(table, "key")
        assert result.num_rows == 2
        pylist = result.to_pylist()
        statuses = {r["key"]: r["status"] for r in pylist}
        assert statuses["A"] == "complete"
        assert statuses["B"] == "failed"


# ── Store filesystem ─────────────────────────────────────────


class TestStoreFilesystem:
    def test_local_path(self, tmp_path) -> None:
        store = Store(tmp_path / "store")
        fs = store._fs()
        assert fs is not None
        # Cached on second call.
        assert store._fs() is fs

    def test_full(self, tmp_path) -> None:
        store = Store(tmp_path / "store")
        assert store._full("a", "b").endswith("/a/b")

    def test_ensure_dir(self, tmp_path) -> None:
        store = Store(tmp_path / "store")
        store._ensure_dir(str(tmp_path / "store" / "nested"))
        assert (tmp_path / "store" / "nested").exists()


# ── Parquet I/O ──────────────────────────────────────────────


class TestParquetIO:
    def test_read_nonexistent(self, tmp_path) -> None:
        store = Store(tmp_path / "store")
        table = store._read_parquet(
            str(tmp_path / "store" / "missing.parquet"),
            JOBS_SCHEMA,
        )
        assert table.num_rows == 0
        assert table.schema.equals(JOBS_SCHEMA)

    def test_read_existing(self, tmp_path) -> None:
        store = Store(tmp_path / "store")
        path = str(tmp_path / "store" / "test.parquet")
        (tmp_path / "store").mkdir()
        table = JOBS_SCHEMA.empty_table()
        with open(path, "wb") as f:
            pq.write_table(table, f)
        result = store._read_parquet(path, JOBS_SCHEMA)
        assert result.num_rows == 0

    def test_append_parquet_creates(self, tmp_path) -> None:
        store = Store(tmp_path / "store")
        path = str(tmp_path / "store" / "jobs.parquet")
        store._append_parquet(
            path,
            JOBS_SCHEMA,
            [
                {
                    "ts": _now(),
                    "key": "X",
                    "job_url": "",
                    "endpoint": "Group/all",
                    "types": "",
                    "since": "",
                    "status": "pending",
                    "error": "",
                }
            ],
        )
        table = pq.read_table(path, schema=JOBS_SCHEMA)
        assert table.num_rows == 1

    def test_append_parquet_appends(self, tmp_path) -> None:
        store = Store(tmp_path / "store")
        path = str(tmp_path / "store" / "jobs.parquet")
        row = {
            "ts": _now(),
            "key": "Y",
            "job_url": "",
            "endpoint": "",
            "types": "",
            "since": "",
            "status": "pending",
            "error": "",
        }
        store._append_parquet(path, JOBS_SCHEMA, [row])
        store._append_parquet(path, JOBS_SCHEMA, [row])
        table = pq.read_table(path, schema=JOBS_SCHEMA)
        assert table.num_rows == 2


class TestAppendJobFile:
    def test_append_job(self, tmp_path) -> None:
        store = Store(tmp_path / "store")
        store._append_job(key="K1", status="pending")
        table = pq.read_table(
            str(tmp_path / "store" / "jobs.parquet"),
            schema=JOBS_SCHEMA,
        )
        assert table.num_rows == 1
        row = table.to_pylist()[0]
        assert row["key"] == "K1"
        assert row["status"] == "pending"
        # Defaults fill in missing columns.
        assert row["job_url"] == ""

    def test_append_file(self, tmp_path) -> None:
        store = Store(tmp_path / "store")
        store._append_file(
            job_key="K1",
            resource_type="Patient",
            url="https://x/data",
            status="pending",
        )
        table = pq.read_table(
            str(tmp_path / "store" / "files.parquet"),
            schema=FILES_SCHEMA,
        )
        assert table.num_rows == 1
        row = table.to_pylist()[0]
        assert row["resource_type"] == "Patient"
        assert row["file_size"] == 0


# ── Query state ──────────────────────────────────────────────


class TestQueryState:
    def test_get_job(self, tmp_path) -> None:
        store = Store(tmp_path / "store")
        store._append_job(key="K1", status="pending")
        store._append_job(key="K1", status="complete")
        job = store.get_job("K1")
        assert job["status"] == "complete"

    def test_get_job_not_found(self, tmp_path) -> None:
        store = Store(tmp_path / "store")
        with pytest.raises(KeyError, match="Job not found"):
            store.get_job("NOPE")

    def test_list_jobs_all(self, tmp_path) -> None:
        store = Store(tmp_path / "store")
        store._append_job(key="A", status="complete")
        store._append_job(key="B", status="failed")
        jobs = store.list_jobs()
        assert len(jobs) == 2

    def test_list_jobs_filtered(self, tmp_path) -> None:
        store = Store(tmp_path / "store")
        store._append_job(key="A", status="complete")
        store._append_job(key="B", status="failed")
        jobs = store.list_jobs(status="failed")
        assert len(jobs) == 1
        assert jobs[0]["key"] == "B"

    def test_get_files(self, tmp_path) -> None:
        store = Store(tmp_path / "store")
        store._append_file(
            job_key="K1",
            url="https://x/1",
            resource_type="Patient",
            status="pending",
        )
        store._append_file(
            job_key="K1",
            url="https://x/1",
            resource_type="Patient",
            status="complete",
        )
        store._append_file(
            job_key="K2",
            url="https://x/2",
            resource_type="EOB",
            status="pending",
        )
        files = store.get_files("K1")
        assert len(files) == 1
        assert files[0]["status"] == "complete"


# ── File storage ─────────────────────────────────────────────


class TestStoreFile:
    def test_store_file(self, tmp_path) -> None:
        store = Store(tmp_path / "store")
        local = tmp_path / "local.ndjson"
        local.write_text('{"id":"1"}\n')
        path = store._store_file(local, "K1", "Patient.ndjson")
        assert "exports/K1/Patient.ndjson" in path
        # Verify the file exists in storage.
        assert store._fs().exists(path)

    def test_open(self, tmp_path) -> None:
        store = Store(tmp_path / "store")
        local = tmp_path / "local.ndjson"
        local.write_text('{"id":"1"}\n')
        path = store._store_file(local, "K1", "Patient.ndjson")
        with store.open(path) as f:
            content = f.read()
        assert b'{"id":"1"}' in content


# ── Export with state tracking ───────────────────────────────


class TestExport:
    def test_success(self, tmp_path) -> None:
        store = Store(tmp_path / "store")
        client = MagicMock()
        client.start_export.return_value = "https://x/jobs/1"
        client.poll.return_value = {
            "output": [
                {"url": "https://x/data/Patient", "type": "Patient"},
            ],
        }

        # _download_file creates a file at the destination path.
        def fake_download(url, dest):
            dest.write_text('{"id":"1"}\n')

        client._download_file.side_effect = fake_download

        key = store.export(
            client,
            types=["Patient"],
            since="2025-01-01",
        )
        assert len(key) == 8

        job = store.get_job(key)
        assert job["status"] == "complete"

        files = store.get_files(key)
        assert len(files) == 1
        assert files[0]["status"] == "complete"

    def test_start_export_failure(self, tmp_path) -> None:
        store = Store(tmp_path / "store")
        client = MagicMock()
        client.start_export.side_effect = RuntimeError("network error")

        with pytest.raises(RuntimeError, match="network error"):
            store.export(client)

        jobs = store.list_jobs(status="failed")
        assert len(jobs) == 1

    def test_poll_failure(self, tmp_path) -> None:
        store = Store(tmp_path / "store")
        client = MagicMock()
        client.start_export.return_value = "https://x/jobs/1"
        client.poll.side_effect = RuntimeError("timeout")

        with pytest.raises(RuntimeError, match="timeout"):
            store.export(client)

        jobs = store.list_jobs(status="failed")
        assert len(jobs) == 1

    def test_download_failure(self, tmp_path) -> None:
        store = Store(tmp_path / "store")
        client = MagicMock()
        client.start_export.return_value = "https://x/jobs/1"
        client.poll.return_value = {
            "output": [
                {"url": "https://x/data/P", "type": "Patient"},
            ],
        }
        client._download_file.side_effect = RuntimeError("disk full")

        with pytest.raises(RuntimeError, match="disk full"):
            store.export(client)

        files = store.get_files(store.list_jobs(status="failed")[0]["key"])
        assert any(f["status"] == "failed" for f in files)

    def test_no_types(self, tmp_path) -> None:
        """Export with types=None passes empty string."""
        store = Store(tmp_path / "store")
        client = MagicMock()
        client.start_export.return_value = "https://x/jobs/1"
        client.poll.return_value = {"output": []}
        key = store.export(client)
        job = store.get_job(key)
        assert job["types"] == ""

    def test_multiple_outputs(self, tmp_path) -> None:
        """Multiple output files with duplicate types."""
        store = Store(tmp_path / "store")
        client = MagicMock()
        client.start_export.return_value = "https://x/jobs/1"
        client.poll.return_value = {
            "output": [
                {"url": "https://x/1", "type": "EOB"},
                {"url": "https://x/2", "type": "EOB"},
            ],
        }

        def fake_download(url, dest):
            dest.write_text("{}\n")

        client._download_file.side_effect = fake_download

        key = store.export(client)
        files = store.get_files(key)
        assert len(files) == 2
        assert all(f["status"] == "complete" for f in files)

    def test_unknown_type(self, tmp_path) -> None:
        """Output entry with missing type defaults to 'unknown'."""
        store = Store(tmp_path / "store")
        client = MagicMock()
        client.start_export.return_value = "https://x/jobs/1"
        client.poll.return_value = {
            "output": [{"url": "https://x/1"}],
        }

        def fake_download(url, dest):
            dest.write_text("{}\n")

        client._download_file.side_effect = fake_download

        key = store.export(client)
        files = store.get_files(key)
        assert files[0]["resource_type"] == "unknown"


# ── Resume ───────────────────────────────────────────────────


class TestResume:
    def test_resume_pending(self, tmp_path) -> None:
        """Resume from pending status re-starts the export."""
        store = Store(tmp_path / "store")
        store._append_job(
            key="K1",
            endpoint="Group/all",
            types="Patient",
            since="2025-01-01",
            status="pending",
        )

        client = MagicMock()
        client.start_export.return_value = "https://x/jobs/1"
        client.poll.return_value = {"output": []}

        key = store.resume(client, "K1")
        assert key == "K1"
        client.start_export.assert_called_once()
        # Status should now be complete.
        job = store.get_job("K1")
        assert job["status"] == "complete"

    def test_resume_polling(self, tmp_path) -> None:
        """Resume from polling status re-polls."""
        store = Store(tmp_path / "store")
        store._append_job(
            key="K1",
            job_url="https://x/jobs/1",
            endpoint="Group/all",
            types="",
            since="",
            status="polling",
        )

        client = MagicMock()
        client.poll.return_value = {"output": []}

        key = store.resume(client, "K1")
        assert key == "K1"
        client.poll.assert_called_once()

    def test_resume_failed(self, tmp_path) -> None:
        """Resume from failed status re-polls."""
        store = Store(tmp_path / "store")
        store._append_job(
            key="K1",
            job_url="https://x/jobs/1",
            endpoint="Group/all",
            types="",
            since="",
            status="failed",
            error="timeout",
        )

        client = MagicMock()
        client.poll.return_value = {"output": []}

        key = store.resume(client, "K1")
        assert key == "K1"
        job = store.get_job("K1")
        assert job["status"] == "complete"

    def test_resume_failed_poll_fails_again(self, tmp_path) -> None:
        """Resume from failed, but poll fails again."""
        store = Store(tmp_path / "store")
        store._append_job(
            key="K1",
            job_url="https://x/jobs/1",
            endpoint="Group/all",
            types="",
            since="",
            status="failed",
        )

        client = MagicMock()
        client.poll.side_effect = RuntimeError("still down")

        with pytest.raises(RuntimeError, match="still down"):
            store.resume(client, "K1")

        job = store.get_job("K1")
        assert job["status"] == "failed"

    def test_resume_complete_with_pending_files(self, tmp_path) -> None:
        """Resume from complete downloads remaining files."""
        store = Store(tmp_path / "store")
        store._append_job(
            key="K1",
            job_url="https://x/jobs/1",
            endpoint="Group/all",
            types="",
            since="",
            status="complete",
        )
        # One file already complete.
        store._append_file(
            job_key="K1",
            url="https://x/data/1",
            resource_type="Patient",
            status="complete",
            storage_path="/store/exports/K1/Patient.ndjson",
            file_size=100,
        )

        client = MagicMock()
        client.poll.return_value = {
            "output": [
                {"url": "https://x/data/1", "type": "Patient"},
                {"url": "https://x/data/2", "type": "Coverage"},
            ],
        }

        def fake_download(url, dest):
            dest.write_text("{}\n")

        client._download_file.side_effect = fake_download

        key = store.resume(client, "K1")
        assert key == "K1"
        # Should only download the missing file.
        client._download_file.assert_called_once()

    def test_resume_complete_all_done(self, tmp_path) -> None:
        """Resume from complete with all files done is a no-op."""
        store = Store(tmp_path / "store")
        store._append_job(
            key="K1",
            job_url="https://x/jobs/1",
            endpoint="Group/all",
            types="",
            since="",
            status="complete",
        )
        store._append_file(
            job_key="K1",
            url="https://x/data/1",
            resource_type="Patient",
            status="complete",
            storage_path="/store/exports/K1/Patient.ndjson",
            file_size=100,
        )

        client = MagicMock()
        client.poll.return_value = {
            "output": [
                {"url": "https://x/data/1", "type": "Patient"},
            ],
        }

        key = store.resume(client, "K1")
        assert key == "K1"
        client._download_file.assert_not_called()

    def test_resume_pending_with_types_and_since(self, tmp_path) -> None:
        """Resume pending with empty types/since passes None."""
        store = Store(tmp_path / "store")
        store._append_job(
            key="K1",
            endpoint="Patient",
            types="",
            since="",
            status="pending",
        )

        client = MagicMock()
        client.start_export.return_value = "https://x/jobs/1"
        client.poll.return_value = {"output": []}

        store.resume(client, "K1")
        client.start_export.assert_called_once_with(
            endpoint="Patient",
            types=None,
            since=None,
        )

    def test_resume_failed_no_job_url(self, tmp_path) -> None:
        """Resume failed with no job_url raises ValueError."""
        store = Store(tmp_path / "store")
        store._append_job(
            key="K1",
            endpoint="Group/all",
            types="",
            since="",
            status="failed",
            job_url="",
        )

        client = MagicMock()

        with pytest.raises(ValueError, match="no job_url"):
            store.resume(client, "K1")

    def test_resume_pending_with_split_types(self, tmp_path) -> None:
        """Resume pending with comma-separated types splits them."""
        store = Store(tmp_path / "store")
        store._append_job(
            key="K1",
            endpoint="Group/all",
            types="Patient,Coverage",
            since="2025-01-01",
            status="pending",
        )

        client = MagicMock()
        client.start_export.return_value = "https://x/jobs/1"
        client.poll.return_value = {"output": []}

        store.resume(client, "K1")
        client.start_export.assert_called_once_with(
            endpoint="Group/all",
            types=["Patient", "Coverage"],
            since="2025-01-01",
        )


# ── Parquet paths ────────────────────────────────────────────


class TestPaths:
    def test_jobs_path(self, tmp_path) -> None:
        store = Store(tmp_path / "store")
        assert store._jobs_path.endswith("jobs.parquet")

    def test_files_path(self, tmp_path) -> None:
        store = Store(tmp_path / "store")
        assert store._files_path.endswith("files.parquet")


# ── Storage options ──────────────────────────────────────────


class TestStorageOptions:
    def test_default_empty(self) -> None:
        store = Store("data/bcda")
        assert store._storage_options == {}

    def test_custom_options(self) -> None:
        opts = {"key": "val"}
        store = Store("s3://bucket", storage_options=opts)
        assert store._storage_options == {"key": "val"}

    def test_root_trailing_slash(self) -> None:
        store = Store("data/bcda/")
        assert store._root == "data/bcda"


# ── v3: adopt completed jobs, type_filter passthrough ───────────


class TestAdoptAndTypeFilter:
    def _client(self, handler):
        import httpx

        from bcda.client import Client

        c = Client("cid", "csec", output_dir="/tmp/bcda-store-test")
        c._http = httpx.Client(transport=httpx.MockTransport(handler))
        return c

    def test_adopt_downloads_completed_job(self, tmp_path) -> None:
        import httpx

        from bcda.store import Store

        def handler(req: httpx.Request) -> httpx.Response:
            url = str(req.url)
            if "/auth/token" in url:
                return httpx.Response(
                    200, json={"access_token": "t", "expires_in": 1200}
                )
            if url.endswith("/jobs/77"):
                return httpx.Response(
                    200,
                    json={
                        "output": [
                            {"type": "Patient", "url": "https://s/data/77/p.ndjson"},
                            {"type": "Patient", "url": "https://s/data/77/p2.ndjson"},
                        ]
                    },
                )
            return httpx.Response(200, content=b'{"resourceType":"Patient"}\n')

        store = Store(path=str(tmp_path))
        client = self._client(handler)
        key = store.adopt(
            client, "https://s/api/v3/jobs/77", types=["Patient"], poll_interval=0
        )
        job = store.get_job(key)
        assert job["status"] == "complete"
        assert job["job_url"].endswith("/jobs/77")
        assert job["types"] == "Patient"
        files = store.get_files(key)
        assert sorted(Path(f["storage_path"]).name for f in files) == [
            "Patient.1.ndjson",
            "Patient.ndjson",
        ]
        assert all(f["status"] == "complete" for f in files)
        client.close()

    def test_adopt_failure_recorded(self, tmp_path) -> None:
        import httpx
        import pytest

        from bcda.client import BcdaHttpError
        from bcda.store import Store

        def handler(req: httpx.Request) -> httpx.Response:
            if "/auth/token" in str(req.url):
                return httpx.Response(
                    200, json={"access_token": "t", "expires_in": 1200}
                )
            return httpx.Response(
                404,
                json={
                    "resourceType": "OperationOutcome",
                    "issue": [{"diagnostics": "Job has expired"}],
                },
            )

        store = Store(path=str(tmp_path))
        client = self._client(handler)
        with pytest.raises(BcdaHttpError, match="expired"):
            store.adopt(client, "https://s/api/v3/jobs/78", poll_interval=0)
        failed = store.list_jobs(status="failed")
        assert len(failed) == 1 and "expired" in failed[0]["error"]
        client.close()

    def test_export_records_type_filter(self, tmp_path) -> None:
        import httpx

        from bcda.store import Store

        seen: dict = {}

        def handler(req: httpx.Request) -> httpx.Response:
            url = str(req.url)
            if "/auth/token" in url:
                return httpx.Response(
                    200, json={"access_token": "t", "expires_in": 1200}
                )
            if "$export" in url:
                seen["url"] = url
                return httpx.Response(
                    202, headers={"Content-Location": "https://s/api/v3/jobs/1"}
                )
            if url.endswith("/jobs/1"):
                return httpx.Response(200, json={"output": []})
            return httpx.Response(200, content=b"")

        store = Store(path=str(tmp_path))
        client = self._client(handler)
        key = store.export(
            client,
            types=["ExplanationOfBenefit"],
            type_filter="ExplanationOfBenefit?outcome=partial",
            poll_interval=0,
        )
        assert "_typeFilter=" in seen["url"]
        assert (
            store.get_job(key)["types"]
            == "ExplanationOfBenefit|_typeFilter=ExplanationOfBenefit?outcome=partial"
        )
        client.close()

    def test_export_retries_after_server_cancellation(self, tmp_path) -> None:
        import httpx
        import pytest

        from bcda.client import JobCancelledError
        from bcda.store import Store

        polls = {"n": 0}

        def handler(req: httpx.Request) -> httpx.Response:
            url = str(req.url)
            if "/auth/token" in url:
                return httpx.Response(
                    200, json={"access_token": "t", "expires_in": 1200}
                )
            if "$export" in url:
                polls["n"] += 1
                return httpx.Response(
                    202,
                    headers={"Content-Location": f"https://s/api/v3/jobs/{polls['n']}"},
                )
            if url.endswith("/jobs/1") or url.endswith("/jobs/2"):
                return httpx.Response(
                    404,
                    json={
                        "resourceType": "OperationOutcome",
                        "issue": [{"diagnostics": "Job has been cancelled."}],
                    },
                )
            if url.endswith("/jobs/3"):
                return httpx.Response(200, json={"output": []})
            return httpx.Response(200, content=b"")

        store = Store(path=str(tmp_path))
        client = self._client(handler)
        key = store.export(client, types=["Patient"], poll_interval=0, cancel_retries=2)
        job = store.get_job(key)
        assert job["status"] == "complete" and job["job_url"].endswith("/jobs/3")
        history = store._read_parquet(
            store._jobs_path, store._read_parquet.__globals__["JOBS_SCHEMA"]
        )
        statuses = [r["status"] for r in history.to_pylist() if r["key"] == key]
        assert statuses.count("cancelled") == 2 and statuses[-1] == "complete"
        client.close()

        # no retries left → the cancellation propagates and is recorded
        polls["n"] = 0
        client = self._client(handler)
        with pytest.raises(JobCancelledError):
            store.export(client, types=["Patient"], poll_interval=0)
        assert store.list_jobs(status="failed")[-1]["error"].startswith(
            "https://s/api/v3/jobs/1"
        )
        client.close()