Spaces:
Running
Running
File size: 4,335 Bytes
13b1a91 a866fc9 13b1a91 a866fc9 13b1a91 a866fc9 13b1a91 a866fc9 13b1a91 a866fc9 13b1a91 a866fc9 13b1a91 a866fc9 9327c55 a866fc9 9327c55 a866fc9 4461572 | 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 | 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
|