learningcircuit--local-deep-research/tests/journal_quality/test_openalex_common.py

"""Tests for the shared helpers in data_sources/_openalex_common.py.

The partition iterator and manifest validator are consumed by BOTH
``openalex.py`` and ``institutions.py`` — a bug here would break both
downloaders the same way. Cover them directly so a regression can't
hide behind either caller's test suite.
"""

from __future__ import annotations

import gzip
import io
from unittest.mock import MagicMock

import pytest

from local_deep_research.journal_quality.data_sources._openalex_common import (
    iter_partitions,
    validate_manifest_entries,
)


def _gz_lines(lines: list[bytes]) -> bytes:
    buf = io.BytesIO()
    with gzip.GzipFile(fileobj=buf, mode="wb") as gz:
        for line in lines:
            gz.write(line + b"\n")
    return buf.getvalue()


def _response(content: bytes) -> MagicMock:
    resp = MagicMock()
    resp.raise_for_status = MagicMock()
    resp.content = content
    return resp


class TestValidateManifestEntries:
    def test_accepts_allowed_prefix(self):
        entries = [
            {"url": "s3://openalex/data/jsonl/sources/part_0.gz"},
            {"url": "s3://openalex/data/jsonl/sources/part_1.gz"},
        ]
        # Should not raise
        validate_manifest_entries(entries, "test")

    def test_rejects_disallowed_prefix(self):
        entries = [
            {"url": "s3://openalex/data/jsonl/sources/part_0.gz"},
            {"url": "s3://attacker/data/evil.gz"},
        ]
        with pytest.raises(ValueError, match="disallowed URL"):
            validate_manifest_entries(entries, "test")

    def test_rejects_missing_url(self):
        entries = [{"url": ""}]
        with pytest.raises(ValueError, match="disallowed URL"):
            validate_manifest_entries(entries, "test")


class TestIterPartitions:
    def test_yields_records_per_partition(self, tmp_path):
        """Each partition yields (idx, total, records) once."""
        entries = [
            {"url": "s3://openalex/data/jsonl/sources/part_0.gz"},
            {"url": "s3://openalex/data/jsonl/sources/part_1.gz"},
        ]
        parts = [
            _gz_lines([b'{"id": "S1"}', b'{"id": "S2"}']),
            _gz_lines([b'{"id": "S3"}']),
        ]
        safe_get = MagicMock(side_effect=[_response(p) for p in parts])

        result = list(
            iter_partitions(
                entries,
                tmp_path,
                file_prefix="t",
                label="test",
                safe_get=safe_get,
            )
        )

        assert len(result) == 2
        idx0, total0, recs0 = result[0]
        idx1, total1, recs1 = result[1]
        assert (idx0, total0) == (0, 2)
        assert (idx1, total1) == (1, 2)
        assert [r["id"] for r in recs0] == ["S1", "S2"]
        assert [r["id"] for r in recs1] == ["S3"]

    def test_cleans_up_tmp_files(self, tmp_path):
        """Temp partition files must be deleted after each partition."""
        entries = [{"url": "s3://openalex/data/jsonl/sources/part_0.gz"}]
        safe_get = MagicMock(
            return_value=_response(_gz_lines([b'{"id": "S1"}']))
        )

        list(
            iter_partitions(
                entries,
                tmp_path,
                file_prefix="cleanup_test",
                label="test",
                safe_get=safe_get,
            )
        )

        # No stray .cleanup_test_part_*.gz files left behind
        leftover = list(tmp_path.glob(".cleanup_test_part_*.gz"))
        assert leftover == []

    def test_cleans_up_tmp_files_on_exception(self, tmp_path):
        """If safe_get raises, the tmp file is still cleaned up."""
        entries = [{"url": "s3://openalex/data/jsonl/sources/part_0.gz"}]
        safe_get = MagicMock(side_effect=RuntimeError("network down"))

        with pytest.raises(RuntimeError):
            list(
                iter_partitions(
                    entries,
                    tmp_path,
                    file_prefix="err_test",
                    label="test",
                    safe_get=safe_get,
                )
            )

        leftover = list(tmp_path.glob(".err_test_part_*.gz"))
        assert leftover == []

    def test_passes_consume_body_to_safe_get(self, tmp_path):
        """``iter_partitions`` must pass ``consume_body=True`` to safe_get.

        Body-stream transients (``ChunkedEncodingError`` / ``ReadTimeout``
        raised during ``resp.content``) are retried by
        ``safe_get_with_retries`` only when ``consume_body=True``. Without
        this flag the wrapper returns headers-only and a mid-stream
        IncompleteRead aborts the whole multi-partition download — the
        exact failure shape that hit CI on the journal-data-integration
        workflow.
        """
        entries = [{"url": "s3://openalex/data/jsonl/sources/part_0.gz"}]
        safe_get = MagicMock(
            return_value=_response(_gz_lines([b'{"id": "S1"}']))
        )

        list(
            iter_partitions(
                entries,
                tmp_path,
                file_prefix="consume_body_test",
                label="test",
                safe_get=safe_get,
            )
        )

        assert safe_get.call_count == 1
        _, kwargs = safe_get.call_args
        assert kwargs.get("consume_body") is True, (
            f"iter_partitions must call safe_get with consume_body=True; "
            f"got kwargs={kwargs!r}"
        )

    def test_uses_higher_retry_budget_than_safe_get_default(self, tmp_path):
        """Partition fetches must override the generic safe_get retry budget.

        ``safe_get_with_retries`` defaults to 3 retries with (1, 2, 4) s
        backoff — ~7 s of total resilience, sized for small request
        bodies. OpenAlex partitions are MB-sized and a sustained S3
        blip (≥ ~10 s of mid-stream IncompleteReads) trips every retry
        inside the same bad window, exhausting the budget and aborting
        the whole multi-partition pull. The release-gate workflow saw
        this twice in a row on 2026-04-26.

        ``iter_partitions`` is responsible for setting a longer-lived
        retry budget. Assert that the kwargs passed to ``safe_get``
        request more than 3 retries AND a backoff schedule whose total
        sleep exceeds the generic 7 s.
        """
        entries = [{"url": "s3://openalex/data/jsonl/sources/part_0.gz"}]
        safe_get = MagicMock(
            return_value=_response(_gz_lines([b'{"id": "S1"}']))
        )

        list(
            iter_partitions(
                entries,
                tmp_path,
                file_prefix="retry_budget_test",
                label="test",
                safe_get=safe_get,
            )
        )

        _, kwargs = safe_get.call_args
        assert kwargs.get("max_retries", 3) > 3, (
            f"partition fetch must use more than the safe_get default of "
            f"3 retries; got max_retries={kwargs.get('max_retries')!r}"
        )
        backoff = kwargs.get("backoff_times")
        assert backoff is not None and sum(backoff) > 7, (
            f"partition fetch must use a longer total backoff than the "
            f"safe_get default of 1+2+4=7s; got backoff_times={backoff!r}"
        )

    def test_forwards_overridden_retry_kwargs_to_safe_get(self, tmp_path):
        """Caller-supplied ``max_retries`` / ``backoff_times`` win.

        Default tuning lives in ``iter_partitions``, but downstream
        callers must still be able to override (e.g. tests that don't
        want to wait, or future callers with different SLAs).
        """
        entries = [{"url": "s3://openalex/data/jsonl/sources/part_0.gz"}]
        safe_get = MagicMock(
            return_value=_response(_gz_lines([b'{"id": "S1"}']))
        )

        list(
            iter_partitions(
                entries,
                tmp_path,
                file_prefix="override_test",
                label="test",
                safe_get=safe_get,
                max_retries=1,
                backoff_times=(0,),
            )
        )

        _, kwargs = safe_get.call_args
        assert kwargs.get("max_retries") == 1
        assert kwargs.get("backoff_times") == (0,)

    def test_suppresses_malformed_lines(self, tmp_path, caplog):
        """Bad JSON lines are skipped, not fatal; first-10 warnings logged."""
        entries = [{"url": "s3://openalex/data/jsonl/sources/part_0.gz"}]
        # One good line + 15 bad lines — the helper should yield the
        # good one, skip the rest, and emit <= 11 warnings (10 per-line
        # + 1 suppression notice, plus 1 aggregate summary).
        lines = [b'{"id": "GOOD"}'] + [b"garbage {{"] * 15
        safe_get = MagicMock(return_value=_response(_gz_lines(lines)))

        result = list(
            iter_partitions(
                entries,
                tmp_path,
                file_prefix="malformed_test",
                label="test",
                safe_get=safe_get,
            )
        )

        assert len(result) == 1
        _, _, recs = result[0]
        assert [r["id"] for r in recs] == ["GOOD"]