|
| 1 | +import csv |
| 2 | +import subprocess |
| 3 | +import tempfile |
| 4 | +import unittest |
| 5 | +from pathlib import Path |
| 6 | + |
| 7 | +from test.helpers import env_with_path, make_executable, seed_conversion_metrics_row |
| 8 | + |
| 9 | + |
| 10 | +REPO_ROOT = Path(__file__).resolve().parents[1] |
| 11 | +SCRIPT = REPO_ROOT / "src" / "compression.sh" |
| 12 | + |
| 13 | + |
| 14 | +def prepare_fake_tools(bin_dir: Path): |
| 15 | + make_executable( |
| 16 | + bin_dir / "gzip", |
| 17 | + """#!/usr/bin/env bash |
| 18 | +set -euo pipefail |
| 19 | +file="${@: -1}" |
| 20 | +cp "$file" "$file.gz" |
| 21 | +""", |
| 22 | + ) |
| 23 | + make_executable( |
| 24 | + bin_dir / "brotli", |
| 25 | + """#!/usr/bin/env bash |
| 26 | +set -euo pipefail |
| 27 | +file="${@: -1}" |
| 28 | +cp "$file" "$file.br" |
| 29 | +""", |
| 30 | + ) |
| 31 | + hdt = bin_dir / "rdf2hdt.sh" |
| 32 | + make_executable( |
| 33 | + hdt, |
| 34 | + """#!/usr/bin/env bash |
| 35 | +set -euo pipefail |
| 36 | +cp "$1" "$2" |
| 37 | +""", |
| 38 | + ) |
| 39 | + return hdt |
| 40 | + |
| 41 | + |
| 42 | +def read_metrics_row(metrics_csv: Path, run_id: str, output_name: str): |
| 43 | + with metrics_csv.open() as f: |
| 44 | + rows = list(csv.DictReader(f)) |
| 45 | + for row in rows: |
| 46 | + if row["run_id"] == run_id and row["output_name"] == output_name: |
| 47 | + return row |
| 48 | + raise AssertionError(f"Metrics row not found for run_id={run_id}, output_name={output_name}") |
| 49 | + |
| 50 | + |
| 51 | +class CompressionUnitTests(unittest.TestCase): |
| 52 | + def test_compression_updates_existing_metrics_row_with_mocked_tools(self): |
| 53 | + with tempfile.TemporaryDirectory() as td: |
| 54 | + tmp_path = Path(td) |
| 55 | + out_root = tmp_path / "out" |
| 56 | + output = out_root / "rdf" |
| 57 | + output.mkdir(parents=True) |
| 58 | + (output / "chunk-a.nq").write_text("<s> <p> <o> .\n") |
| 59 | + (output / "chunk-b.nq").write_text("<s2> <p2> <o2> .\n") |
| 60 | + |
| 61 | + logdir = tmp_path / "metrics" |
| 62 | + run_id = "run-compress-1" |
| 63 | + timestamp = "2026-01-01T00:00:00" |
| 64 | + metrics_csv = logdir / "metrics.csv" |
| 65 | + seed_conversion_metrics_row(metrics_csv, run_id, timestamp, "rdf", output) |
| 66 | + |
| 67 | + fake_bin = tmp_path / "bin" |
| 68 | + fake_bin.mkdir() |
| 69 | + hdt_path = prepare_fake_tools(fake_bin) |
| 70 | + |
| 71 | + env = env_with_path(fake_bin) |
| 72 | + env.update( |
| 73 | + { |
| 74 | + "OUT_ROOT_DIR": str(out_root), |
| 75 | + "OUT_NAME": "rdf", |
| 76 | + "LOGDIR": str(logdir), |
| 77 | + "RUN_ID": run_id, |
| 78 | + "TIMESTAMP": timestamp, |
| 79 | + "RDF2HDT": str(hdt_path), |
| 80 | + } |
| 81 | + ) |
| 82 | + |
| 83 | + result = subprocess.run( |
| 84 | + ["bash", str(SCRIPT), "-m", "gzip,brotli,hdt"], |
| 85 | + env=env, |
| 86 | + capture_output=True, |
| 87 | + text=True, |
| 88 | + ) |
| 89 | + |
| 90 | + self.assertEqual(result.returncode, 0, msg=result.stderr) |
| 91 | + self.assertTrue((output / "rdf.nq.gz").exists()) |
| 92 | + self.assertTrue((output / "rdf.nq.br").exists()) |
| 93 | + self.assertTrue((output / "rdf.hdt").exists()) |
| 94 | + |
| 95 | + row = read_metrics_row(metrics_csv, run_id, "rdf") |
| 96 | + self.assertEqual(row["run_id"], run_id) |
| 97 | + self.assertEqual(row["output_name"], "rdf") |
| 98 | + self.assertEqual(row["exit_code_java"], "0") |
| 99 | + self.assertEqual(row["compression_methods"], "gzip|brotli|hdt") |
| 100 | + self.assertEqual(row["exit_code_gzip"], "0") |
| 101 | + self.assertEqual(row["exit_code_brotli"], "0") |
| 102 | + self.assertEqual(row["exit_code_hdt"], "0") |
| 103 | + self.assertGreater(int(row["combined_nq_size_bytes"]), 0) |
| 104 | + |
| 105 | + def test_compression_none_updates_metrics_without_generating_outputs(self): |
| 106 | + with tempfile.TemporaryDirectory() as td: |
| 107 | + tmp_path = Path(td) |
| 108 | + out_root = tmp_path / "out" |
| 109 | + output = out_root / "rdf" |
| 110 | + output.mkdir(parents=True) |
| 111 | + (output / "chunk-a.nq").write_text("<s> <p> <o> .\n") |
| 112 | + |
| 113 | + logdir = tmp_path / "metrics" |
| 114 | + run_id = "run-compress-2" |
| 115 | + timestamp = "2026-01-01T00:00:00" |
| 116 | + metrics_csv = logdir / "metrics.csv" |
| 117 | + seed_conversion_metrics_row(metrics_csv, run_id, timestamp, "rdf", output) |
| 118 | + |
| 119 | + fake_bin = tmp_path / "bin" |
| 120 | + fake_bin.mkdir() |
| 121 | + hdt_path = prepare_fake_tools(fake_bin) |
| 122 | + |
| 123 | + env = env_with_path(fake_bin) |
| 124 | + env.update( |
| 125 | + { |
| 126 | + "OUT_ROOT_DIR": str(out_root), |
| 127 | + "OUT_NAME": "rdf", |
| 128 | + "LOGDIR": str(logdir), |
| 129 | + "RUN_ID": run_id, |
| 130 | + "TIMESTAMP": timestamp, |
| 131 | + "RDF2HDT": str(hdt_path), |
| 132 | + } |
| 133 | + ) |
| 134 | + |
| 135 | + result = subprocess.run( |
| 136 | + ["bash", str(SCRIPT), "-m", "none"], |
| 137 | + env=env, |
| 138 | + capture_output=True, |
| 139 | + text=True, |
| 140 | + ) |
| 141 | + |
| 142 | + self.assertEqual(result.returncode, 0, msg=result.stderr) |
| 143 | + self.assertFalse((output / "rdf.nq.gz").exists()) |
| 144 | + self.assertFalse((output / "rdf.nq.br").exists()) |
| 145 | + self.assertFalse((output / "rdf.hdt").exists()) |
| 146 | + |
| 147 | + row = read_metrics_row(metrics_csv, run_id, "rdf") |
| 148 | + self.assertEqual(row["compression_methods"], "none") |
| 149 | + self.assertEqual(row["combined_nq_size_bytes"], "0") |
| 150 | + self.assertEqual(row["gzip_size_bytes"], "0") |
| 151 | + self.assertEqual(row["brotli_size_bytes"], "0") |
| 152 | + self.assertEqual(row["hdt_size_bytes"], "0") |
| 153 | + |
| 154 | + |
| 155 | +if __name__ == "__main__": |
| 156 | + unittest.main() |
0 commit comments