stack/tests/bib/test_regulations_gov_mirror.py

"""Tests for the Mirrulations S3 mirror backfill (bib/regulations_gov.py)."""

from __future__ import annotations

import json

import httpx

from bib.regulations_gov import (
    Comment,
    Mirror,
    backfill_from_mirror,
    upsert_comment,
)
from bib.store import Store

_DOCKET = "CMS-2018-0076"
_TEXT = f"raw-data/CMS/{_DOCKET}/text-{_DOCKET}/comments/"
_BIN = f"raw-data/CMS/{_DOCKET}/binary-{_DOCKET}/comments_attachments/"


def _detail(cid: str, comment: str, org: str = "") -> dict:
    return {
        "data": {
            "id": cid,
            "type": "comments",
            "attributes": {
                "title": f"Comment {cid}",
                "docketId": _DOCKET,
                "commentOnId": "0900006482921ba1",
                "postedDate": "2018-09-10T00:00:00Z",
                "receivedDate": "2018-09-09T00:00:00Z",
                "organization": org,
                "comment": comment,
                "attachmentCount": 1,
            },
        }
    }


_PDF = b"%PDF-1.4 " + b"x" * 2000


def _handler(request: httpx.Request) -> httpx.Response:
    path = request.url.path.lstrip("/")
    params = dict(request.url.params)
    if path == "" and "list-type" in params:
        prefix = params.get("prefix", "")
        if prefix == _TEXT:
            if "continuation-token" not in params:
                return httpx.Response(
                    200,
                    text=(
                        "<ListBucketResult>"
                        f"<Contents><Key>{_TEXT}{_DOCKET}-0001.json</Key></Contents>"
                        "<NextContinuationToken>tok1</NextContinuationToken>"
                        "</ListBucketResult>"
                    ),
                )
            return httpx.Response(
                200,
                text=(
                    "<ListBucketResult>"
                    f"<Contents><Key>{_TEXT}{_DOCKET}-0002.json</Key></Contents>"
                    "</ListBucketResult>"
                ),
            )
        if prefix == _BIN:
            return httpx.Response(
                200,
                text=(
                    "<ListBucketResult>"
                    f"<Contents><Key>{_BIN}{_DOCKET}-0002_attachment_1.pdf</Key>"
                    "</Contents></ListBucketResult>"
                ),
            )
        return httpx.Response(200, text="<ListBucketResult></ListBucketResult>")
    if path == f"{_TEXT}{_DOCKET}-0001.json":
        return httpx.Response(
            200, text=json.dumps(_detail(f"{_DOCKET}-0001", "inline body one"))
        )
    if path == f"{_TEXT}{_DOCKET}-0002.json":
        return httpx.Response(
            200,
            text=json.dumps(_detail(f"{_DOCKET}-0002", "see attached", org="Test Org")),
        )
    if path == f"{_BIN}{_DOCKET}-0002_attachment_1.pdf":
        return httpx.Response(
            200, content=_PDF, headers={"content-type": "application/pdf"}
        )
    return httpx.Response(404)


def _mirror() -> Mirror:
    return Mirror(http=httpx.Client(transport=httpx.MockTransport(_handler)))


class TestMirrorListing:
    def test_paged_listing_concatenates(self):
        ids = _mirror().comment_ids(_DOCKET)
        assert ids == [f"{_DOCKET}-0001", f"{_DOCKET}-0002"]

    def test_attachment_keys_by_cid(self):
        m = _mirror().attachment_keys(_DOCKET)
        assert list(m) == [f"{_DOCKET}-0002"]
        assert m[f"{_DOCKET}-0002"] == [f"{_BIN}{_DOCKET}-0002_attachment_1.pdf"]


class TestBackfillFromMirror:
    def _store(self, tmp_path) -> Store:
        store = Store(str(tmp_path / "bib.sqlite"))
        # Seed one un-enriched stub the way the original farm did.
        c = Comment(
            id=f"{_DOCKET}-0001",
            title="",
            posted_date="2018-09-10",
            received_date="2018-09-09",
            docket_id=_DOCKET,
            comment_on_id="x",
        )
        upsert_comment(
            store, c, cms_id="CMS-1693-P", extra_tags=[f"reg-docket:{_DOCKET}"]
        )
        return store

    def test_enriches_creates_and_attaches(self, tmp_path):
        store = self._store(tmp_path)
        stats = backfill_from_mirror(
            store,
            _mirror(),
            docket=_DOCKET,
            scratch_root=tmp_path / "scratch",
            workers=2,
        )
        assert stats["enriched"] == 2
        assert stats["created"] == 1
        assert stats["attached"] == 1
        assert stats["errors"] == 0

        con = store._con()
        # stub got its body + enriched:ok
        row = con.execute(
            "SELECT abstract FROM items WHERE url LIKE ?",
            (f"%{_DOCKET}-0001",),
        ).fetchone()
        assert row[0] == "inline body one"
        # mirror-only comment created with the farm tag set
        row2 = con.execute(
            "SELECT id FROM items WHERE url LIKE ?", (f"%{_DOCKET}-0002",)
        ).fetchone()
        assert row2 is not None
        tags = {
            r[0]
            for r in con.execute(
                "SELECT t.name FROM tags t JOIN item_tags it ON it.tag_id=t.id "
                "WHERE it.item_id=?",
                (row2[0],),
            )
        }
        assert "source:regulations-gov" in tags
        assert f"reg-docket:{_DOCKET}" in tags
        assert "rule:CMS-1693-P" in tags
        assert "enriched:ok" in tags
        assert "org:test-org" in tags
        # attachment landed with the layout name
        att = tmp_path / "scratch" / _DOCKET / f"{_DOCKET}-0002" / "attachment_1.pdf"
        assert att.is_file() and att.read_bytes().startswith(b"%PDF")

    def test_idempotent_rerun(self, tmp_path):
        store = self._store(tmp_path)
        m = _mirror()
        backfill_from_mirror(
            store, m, docket=_DOCKET, scratch_root=tmp_path / "scratch", workers=2
        )
        stats2 = backfill_from_mirror(
            store, m, docket=_DOCKET, scratch_root=tmp_path / "scratch", workers=2
        )
        assert stats2["enriched"] == 0
        assert stats2["created"] == 0