Spaces:
Running
Running
| import pyarrow.parquet as pq | |
| import pytest | |
| from app.events import COLUMNS, ClientEvent, to_row | |
| from app.sink import HubParquetSink, LocalParquetSink, rows_to_parquet_bytes | |
| def rows(n): | |
| return [to_row(ClientEvent(event_type="session", model_family="Qwen", quantization="4-bit")) for _ in range(n)] | |
| def test_batch_becomes_one_shard(tmp_path): | |
| s = LocalParquetSink(tmp_path, flush_seconds=999) | |
| s.add(rows(7)) | |
| s.add(rows(5)) | |
| assert s.flush() | |
| shards = list((tmp_path / "data/incoming").rglob("*.parquet")) | |
| assert len(shards) == 1 | |
| t = pq.read_table(shards[0]) | |
| assert t.num_rows == 12 and t.column_names == list(COLUMNS) | |
| assert s.flush() and len(list((tmp_path / "data/incoming").rglob("*.parquet"))) == 1 # empty flush = no shard | |
| s.add(rows(1)) | |
| s.flush() | |
| assert len(list((tmp_path / "data/incoming").rglob("*.parquet"))) == 2 # appends, never rewrites | |
| assert len(s.read_existing()) == 13 | |
| class FlakyApi: | |
| def __init__(self, fail_times): | |
| self.fail_times = fail_times | |
| self.commits = [] | |
| self.token = None | |
| def create_commit(self, repo_id, repo_type, operations, commit_message): | |
| if self.fail_times > 0: | |
| self.fail_times -= 1 | |
| raise ConnectionError("hub down") | |
| assert repo_type == "dataset" | |
| self.commits.append([op.path_in_repo for op in operations]) | |
| def test_hub_failure_keeps_spool_and_retries(tmp_path): | |
| api = FlakyApi(fail_times=2) | |
| s = HubParquetSink("org/data", None, spool_dir=tmp_path, api=api, flush_seconds=999) | |
| s.add(rows(3)) | |
| assert not s.flush() and s.pending == 3 and not s.status()["healthy"] | |
| # a restart in between must not lose rows | |
| s2 = HubParquetSink("org/data", None, spool_dir=tmp_path, api=api, flush_seconds=999) | |
| assert s2.pending == 3 | |
| assert not s2.flush() | |
| assert s2.flush() and s2.pending == 0 and s2.status()["healthy"] | |
| assert len(api.commits) == 1 and len(api.commits[0]) == 1 | |
| assert api.commits[0][0].startswith("data/incoming/") and api.commits[0][0].endswith(".parquet") | |
| s3 = HubParquetSink("org/data", None, spool_dir=tmp_path, api=api, flush_seconds=999) | |
| assert s3.pending == 0 | |
| def test_local_refuses_overwrite(tmp_path): | |
| s = LocalParquetSink(tmp_path, flush_seconds=999) | |
| data = rows_to_parquet_bytes(rows(1)) | |
| s._write_shard("data/incoming/x.parquet", data) | |
| with pytest.raises(FileExistsError): | |
| s._write_shard("data/incoming/x.parquet", data) | |
| def test_dataset_schema_is_deliberate(): | |
| """Changing columns changes the public dataset schema: bump this list on purpose, and | |
| migrate or document existing shards (see the dataset card's versioning section).""" | |
| assert len(COLUMNS) == 67 | |
| assert {"perplexity", "perplexity_stderr", "eval_dataset", "eval_tokens"} <= set(COLUMNS) | |
| def test_size_triggered_flush(tmp_path, monkeypatch): | |
| import time | |
| import app.sink as sink_mod | |
| monkeypatch.setattr(sink_mod, "FLUSH_ROWS", 5) | |
| s = LocalParquetSink(tmp_path, flush_seconds=999) | |
| s.add(rows(3)) | |
| assert s.pending == 3 | |
| s.add(rows(3)) # crosses FLUSH_ROWS -> flushes in the background | |
| for _ in range(50): | |
| if s.pending == 0: | |
| break | |
| time.sleep(0.05) | |
| assert s.pending == 0 and len(list((tmp_path / "data/incoming").rglob("*.parquet"))) == 1 | |
| def test_reads_monthly_and_incoming(tmp_path): | |
| from app.sink import rows_to_parquet_bytes as rb | |
| (tmp_path / "data/events").mkdir(parents=True) | |
| (tmp_path / "data/events/2026-08.parquet").write_bytes(rb(rows(4))) | |
| s = LocalParquetSink(tmp_path, flush_seconds=999) | |
| s.add(rows(2)) | |
| s.flush() | |
| assert len(s.read_existing()) == 6 | |
| def test_reload_neither_misses_nor_double_counts(tmp_path): | |
| import time | |
| s = LocalParquetSink(tmp_path, flush_seconds=999) | |
| s.add(rows(3)) | |
| stored = s.read_existing() | |
| assert len(stored) + len(s.live_rows()) == 3 # accepted, not yet written | |
| s.flush() | |
| assert len(s.read_existing()) + len(s.live_rows()) == 6 # written but not yet reloaded: still visible (overlap is fine before a reload) | |
| started = time.time() | |
| stored = s.read_existing() | |
| s.forget_flushed_before(started) | |
| s.add(rows(2)) | |
| assert len(stored) + len(s.live_rows()) == 5 # after reload: 3 stored + 2 new, no duplicates | |