stack/tests/zot/test_duck.py

"""Tests for zot.duck — DuckDB engine for Zotero analytics."""

from __future__ import annotations

from pathlib import Path

import pytest

from zot.db import TYPE_MAP, Db

# We need a populated SQLite DB for DuckDB to attach
HOST_DB = Path("data/zotero/data/zotero.sqlite")


@pytest.fixture()
def sqlite_db(zotero_db) -> str:
    """Populated Zotero SQLite for testing — uses the session-scoped
    template fixture (~10s saved per test by skipping create_db)."""
    path = zotero_db
    with Db(path) as db:
        # Create a few items
        id1 = db.create_item(TYPE_MAP["statute"])
        db.set_fields(
            id1, {"nameOfAct": "PFS 2026", "code": "FR", "url": "https://ex.com/pfs"}
        )
        db.sync_tags(id1, ["module:pfs", "year:2026"])

        id2 = db.create_item(TYPE_MAP["journalArticle"])
        db.set_fields(
            id2,
            {
                "title": "Skin Substitutes Review",
                "DOI": "10.1234/test",
                "publicationTitle": "JAMA",
                "date": "2025-06-01",
            },
        )
        db.add_creators(id2, [("Jane", "Smith"), ("Bob", "Lee")])
        db.sync_tags(id2, ["module:skin-subs", "year:2025"])

        id3 = db.create_item(TYPE_MAP["webpage"])
        db.set_fields(id3, {"title": "CMS Data", "url": "https://cms.gov/data"})
        db.sync_tags(id3, ["module:pfs"])

        key = db.ensure_collection("Test Collection")
        db.add_to_collection(id1, collection_key=key)
        db.add_to_collection(id2, collection_key=key)

        db.commit()
    return path


class TestAttach:
    def test_attach_and_query(self, sqlite_db):
        from zot.duck import DuckDb

        with DuckDb.attach(sqlite_db) as zdb:
            result = zdb.sql("SELECT count(*) FROM zot.items").fetchone()
            assert result[0] == 3

    def test_items_flat(self, sqlite_db):
        from zot.duck import DuckDb

        with DuckDb.attach(sqlite_db) as zdb:
            rel = zdb.items()
            rows = rel.fetchall()
            assert len(rows) == 3
            # Check columns exist
            cols = [d[0] for d in rel.description]
            assert "title" in cols
            assert "doi" in cols
            assert "item_type" in cols

    def test_tags_flat(self, sqlite_db):
        from zot.duck import DuckDb

        with DuckDb.attach(sqlite_db) as zdb:
            rows = zdb.tags().fetchall()
            tags = {r[2] for r in rows}  # tag column
            assert "module:pfs" in tags
            assert "module:skin-subs" in tags

    def test_creators_flat(self, sqlite_db):
        from zot.duck import DuckDb

        with DuckDb.attach(sqlite_db) as zdb:
            rows = zdb.creators().fetchall()
            names = {r[2] for r in rows}  # last_name column
            assert "Smith" in names
            assert "Lee" in names

    def test_collections_flat(self, sqlite_db):
        from zot.duck import DuckDb

        with DuckDb.attach(sqlite_db) as zdb:
            rows = zdb.collections().fetchall()
            assert len(rows) == 2  # two items in collection

    def test_items_by_type(self, sqlite_db):
        from zot.duck import DuckDb

        with DuckDb.attach(sqlite_db) as zdb:
            rows = zdb.items_by_type().fetchall()
            types = {r[0]: r[1] for r in rows}
            assert types.get("statute") == 1
            assert types.get("journalArticle") == 1
            assert types.get("webpage") == 1

    def test_items_by_tag(self, sqlite_db):
        from zot.duck import DuckDb

        with DuckDb.attach(sqlite_db) as zdb:
            rows = zdb.items_by_tag().fetchall()
            tags = {r[0]: r[1] for r in rows}
            assert tags.get("module:pfs") == 2

    def test_doi_coverage(self, sqlite_db):
        from zot.duck import DuckDb

        with DuckDb.attach(sqlite_db) as zdb:
            rows = zdb.doi_coverage().fetchall()
            assert len(rows) > 0


class TestMirror:
    def test_mirror_creates_native_tables(self, sqlite_db):
        from zot.duck import DuckDb

        with DuckDb.mirror(sqlite_db) as zdb:
            count = zdb.sql("SELECT count(*) FROM items").fetchone()[0]
            assert count == 3

    def test_mirror_materialize(self, sqlite_db):
        from zot.duck import DuckDb

        with DuckDb.mirror(sqlite_db) as zdb:
            zdb.materialize()
            rows = zdb.sql("SELECT count(*) FROM flat_items").fetchone()
            assert rows[0] == 3

    def test_mirror_to_parquet(self, sqlite_db, tmp_path):
        from zot.duck import DuckDb

        with DuckDb.mirror(sqlite_db) as zdb:
            zdb.materialize()
            parquet_path = str(tmp_path / "items.parquet")
            zdb.to_parquet(parquet_path)
            assert Path(parquet_path).exists()
            assert Path(parquet_path).stat().st_size > 0


class TestToDataFrame:
    def test_to_polars(self, sqlite_db):
        from zot.duck import DuckDb

        with DuckDb.attach(sqlite_db) as zdb:
            df = zdb.to_df()
            assert len(df) == 3
            assert "title" in df.columns
            assert "doi" in df.columns


class TestFromDb:
    """Line 186: DuckDb.from_db constructs from an open Db instance."""

    def test_from_db(self, sqlite_db):
        from zot.duck import DuckDb

        with Db(sqlite_db) as db:
            with DuckDb.from_db(db) as zdb:
                count = zdb.sql("SELECT count(*) FROM items").fetchone()[0]
                assert count == 3


class TestSqlWithParams:
    """Line 193: sql() with params list."""

    def test_sql_with_params(self, sqlite_db):
        from zot.duck import DuckDb

        with DuckDb.attach(sqlite_db) as zdb:
            result = zdb.sql(
                "SELECT count(*) FROM zot.items WHERE itemTypeID = ?",
                [TYPE_MAP["statute"]],
            ).fetchone()
            assert result[0] == 1


class TestItemsByYear:
    """Line 257: items_by_year analytics query."""

    def test_items_by_year(self, sqlite_db):
        from zot.duck import DuckDb

        with DuckDb.attach(sqlite_db) as zdb:
            rows = zdb.items_by_year().fetchall()
            assert len(rows) >= 1


class TestItemsByCollection:
    """Line 267: items_by_collection analytics query."""

    def test_items_by_collection(self, sqlite_db):
        from zot.duck import DuckDb

        with DuckDb.attach(sqlite_db) as zdb:
            rows = zdb.items_by_collection().fetchall()
            assert len(rows) >= 1
            # Should have our "Test Collection"
            names = {r[0] for r in rows}
            assert "Test Collection" in names


@pytest.mark.skipif(not HOST_DB.exists(), reason="zotero.sqlite not available")
class TestRealDb:
    """Integration tests against a snapshot of the real Zotero database.

    Uses the session-scoped ``host_db`` fixture (see conftest.py) which
    snapshots the live DB once — opening the live file directly fails
    with ``database is locked`` while the Zotero container holds WAL.
    """

    def test_attach_real_db(self, host_db):
        from zot.duck import DuckDb

        with DuckDb.attach(str(host_db)) as zdb:
            count = zdb.sql("SELECT count(*) FROM zot.items").fetchone()[0]
            assert count > 1000

    def test_flat_items_real(self, host_db):
        from zot.duck import DuckDb

        with DuckDb.attach(str(host_db)) as zdb:
            df = zdb.to_df()
            assert len(df) > 100
            assert "doi" in df.columns

    def test_tag_co_occurrence_real(self, host_db):
        from zot.duck import DuckDb

        with DuckDb.attach(str(host_db)) as zdb:
            rows = zdb.tag_co_occurrence(min_count=10).fetchall()
            assert len(rows) > 0

    def test_pdf_coverage_real(self, host_db):
        from zot.duck import DuckDb

        with DuckDb.attach(str(host_db)) as zdb:
            rows = zdb.pdf_coverage().fetchall()
            assert len(rows) > 0