tests/test_index_query.py

"""Exercise the estimate-job payload (src/index_query.py) in local mode.

This is the actual code that runs in the estimate job: it builds the SQL (vendored,
no cdx_toolkit), reads the parquet index with DuckDB, writes the range-jobs CSV, and
prints the ESTIMATE summary. Local mode globs the fixture dir instead of enumerating
CDN URLs, so it runs fully offline.
"""
import os

import pytest

pytest.importorskip("duckdb")


def test_index_query_local(cc_fixtures, tmp_path, monkeypatch, capsys):
    from src import index_query

    out = tmp_path / "ranges.csv"
    monkeypatch.setenv("CC_INDEX_MODE", "local")
    monkeypatch.setenv("CC_LOCAL_INDEX_DIR", os.path.join(cc_fixtures, "cc-index/table/cc-main/warc"))
    monkeypatch.setenv("CC_CRAWLS", "CC-MAIN-2026-25")
    monkeypatch.setenv("CC_DOMAINS", "example.com")
    monkeypatch.setenv("CC_HOSTNAMES", "blog.example.org")
    monkeypatch.setenv("CC_RANGES_OUT", str(out))

    rc = index_query.main()
    assert rc == 0

    captured = capsys.readouterr().out
    assert "ESTIMATE n_records=2 total_bytes=" in captured

    rows = out.read_text().strip().splitlines()
    assert rows[0] == "warc_filename,warc_record_offset,warc_record_length"
    assert len(rows) == 3  # header + 2 matched records (example.com + blog.example.org)
    assert all("crawl-data/CC-MAIN-2026-25/" in r for r in rows[1:])


def _lang_env(cc_fixtures, out, monkeypatch, languages):
    monkeypatch.setenv("CC_INDEX_MODE", "local")
    monkeypatch.setenv("CC_LOCAL_INDEX_DIR", os.path.join(cc_fixtures, "cc-index/table/cc-main/warc"))
    monkeypatch.setenv("CC_CRAWLS", "CC-MAIN-2026-25")
    monkeypatch.setenv("CC_DOMAINS", "example.com,example.org")
    monkeypatch.setenv("CC_HOSTNAMES", "")
    monkeypatch.setenv("CC_LANGUAGES", languages)
    monkeypatch.setenv("CC_RANGES_OUT", str(out))


def test_index_query_language_filter_primary(cc_fixtures, tmp_path, monkeypatch, capsys):
    """deu matches the "deu,eng" record via the LIKE 'deu,%' prefix."""
    from src import index_query

    out = tmp_path / "ranges.csv"
    _lang_env(cc_fixtures, out, monkeypatch, "deu")
    assert index_query.main() == 0
    assert "ESTIMATE n_records=1 total_bytes=" in capsys.readouterr().out


def test_index_query_language_ignores_secondary(cc_fixtures, tmp_path, monkeypatch, capsys):
    """The filter is on the PRIMARY language: eng matches the "eng" record but not the
    "deu,eng" one, where eng is only the secondary detection."""
    from src import index_query

    out = tmp_path / "ranges.csv"
    _lang_env(cc_fixtures, out, monkeypatch, "eng")
    assert index_query.main() == 0
    assert "ESTIMATE n_records=1 total_bytes=" in capsys.readouterr().out


def test_index_query_language_filter_multi(cc_fixtures, tmp_path, monkeypatch, capsys):
    """Several codes are OR-ed: eng ("eng") + deu ("deu,eng") match both records."""
    from src import index_query

    out = tmp_path / "ranges.csv"
    _lang_env(cc_fixtures, out, monkeypatch, "eng,deu")
    assert index_query.main() == 0
    assert "ESTIMATE n_records=2 total_bytes=" in capsys.readouterr().out


def test_index_query_language_no_match(cc_fixtures, tmp_path, monkeypatch, capsys):
    from src import index_query

    out = tmp_path / "ranges.csv"
    _lang_env(cc_fixtures, out, monkeypatch, "jpn")
    assert index_query.main() == 0
    assert "ESTIMATE n_records=0 total_bytes=0" in capsys.readouterr().out


def test_index_query_rejects_bad_language_code(cc_fixtures, tmp_path, monkeypatch, capsys):
    """Codes are inlined as SQL literals, so anything not [a-z]{3} is rejected."""
    from src import index_query

    out = tmp_path / "ranges.csv"
    _lang_env(cc_fixtures, out, monkeypatch, "eng'); DROP TABLE x;--")
    assert index_query.main() == 1
    assert "invalid ISO-639-3 language code" in capsys.readouterr().err


def test_index_query_zero_matches(cc_fixtures, tmp_path, monkeypatch, capsys):
    """A filter that matches no rows must still report ESTIMATE n_records=0 (not crash
    on sum() over a header-only/VARCHAR CSV)."""
    from src import index_query

    out = tmp_path / "ranges.csv"
    monkeypatch.setenv("CC_INDEX_MODE", "local")
    monkeypatch.setenv("CC_LOCAL_INDEX_DIR", os.path.join(cc_fixtures, "cc-index/table/cc-main/warc"))
    monkeypatch.setenv("CC_CRAWLS", "CC-MAIN-2026-25")
    monkeypatch.setenv("CC_DOMAINS", "")
    monkeypatch.setenv("CC_HOSTNAMES", "nope.example.invalid")  # matches nothing
    monkeypatch.setenv("CC_RANGES_OUT", str(out))

    assert index_query.main() == 0
    assert "ESTIMATE n_records=0 total_bytes=0" in capsys.readouterr().out
    assert out.read_text().strip() == "warc_filename,warc_record_offset,warc_record_length"


def test_index_query_wraps_errors(monkeypatch, capsys):
    """Unexpected failures become a clean ERROR line + non-zero exit, not a traceback."""
    from src import index_query

    monkeypatch.setattr(index_query, "_run", lambda: (_ for _ in ()).throw(RuntimeError("boom")))
    assert index_query.main() == 1
    assert "ERROR: index query failed: RuntimeError: boom" in capsys.readouterr().err


def test_index_query_no_files(tmp_path, monkeypatch):
    from src import index_query

    monkeypatch.setenv("CC_INDEX_MODE", "local")
    monkeypatch.setenv("CC_LOCAL_INDEX_DIR", str(tmp_path / "empty"))
    monkeypatch.setenv("CC_CRAWLS", "CC-MAIN-2026-25")
    monkeypatch.setenv("CC_DOMAINS", "example.com")
    monkeypatch.setenv("CC_HOSTNAMES", "")
    monkeypatch.setenv("CC_RANGES_OUT", str(tmp_path / "ranges.csv"))
    assert index_query.main() == 3  # no parquet files found