Spaces:
Sleeping
Sleeping
File size: 5,603 Bytes
2abe74e 8eaca44 2abe74e 8eaca44 d9b42dc 2abe74e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 | """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
|