"""内容指纹:识别"换个包名重传同一份数据"。 实测背景:采集端把已上传的录制重新推一遍时包名会变(本地去重标记丢失 或包名规则变化),``session.py`` 的 incoming_name 判断不出来,同一份数据 被追加成新 episode(UMIGripper_AI 出现 4 组字节级相同的批次)。这里的 指纹必须与包名无关,且对"同一次录制的两份拷贝"稳定命中。 """ import json import tempfile import unittest from pathlib import Path from unittest import mock import numpy as np import pyarrow as pa import pyarrow.parquet as pq from app import batch_fingerprint def _write_episode(root: Path, episode_index: int, video_bytes: bytes, frames: int = 4) -> None: """写一个最小 canonical 布局:一个视频 + 一个 data parquet。""" videos = root / "videos" / "observation.images.gripper_rgb" / "chunk-000" videos.mkdir(parents=True, exist_ok=True) (videos / f"episode_{episode_index:06d}.mp4").write_bytes(video_bytes) data = root / "data" / "chunk-000" data.mkdir(parents=True, exist_ok=True) table = pa.table({ "episode_index": pa.array([episode_index] * frames, pa.int64()), "frame_index": pa.array(list(range(frames)), pa.int64()), "observation.value": pa.array( [[float(i), float(i) + 1.0] for i in range(frames)], pa.list_(pa.float32()), ), }) pq.write_table(table, data / f"episode_{episode_index:06d}.parquet") class FingerprintTests(unittest.TestCase): def setUp(self): self._tmp = tempfile.TemporaryDirectory() self.root = Path(self._tmp.name) # 索引落到临时目录,别碰真实 state/ self._index_patch = mock.patch.object( batch_fingerprint, "INDEX_PATH", self.root / "fingerprints.json") self._index_patch.start() def tearDown(self): self._index_patch.stop() self._tmp.cleanup() def test_same_recording_different_episode_index_matches(self): """同一次录制的两份拷贝,只有 episode_index 不同 → 指纹相同。""" a = self.root / "a" b = self.root / "b" _write_episode(a, 3, b"IDENTICAL-VIDEO-BYTES") _write_episode(b, 4, b"IDENTICAL-VIDEO-BYTES") self.assertIsNotNone(batch_fingerprint.compute_fingerprint(a)) self.assertEqual( batch_fingerprint.compute_fingerprint(a), batch_fingerprint.compute_fingerprint(b), ) def test_different_video_differs(self): a = self.root / "a" b = self.root / "b" _write_episode(a, 0, b"VIDEO-ONE") _write_episode(b, 0, b"VIDEO-TWO") self.assertNotEqual( batch_fingerprint.compute_fingerprint(a), batch_fingerprint.compute_fingerprint(b), ) def test_different_frame_data_differs(self): """视频相同但帧数据不同 → 不同指纹(不能只比视频)。""" a = self.root / "a" b = self.root / "b" _write_episode(a, 0, b"SAME-VIDEO") _write_episode(b, 0, b"SAME-VIDEO") # 只改 b 的一帧数据 path = b / "data" / "chunk-000" / "episode_000000.parquet" table = pq.read_table(path).to_pydict() table["observation.value"] = [[9.0, 9.0] for _ in table["frame_index"]] pq.write_table(pa.table(table), path) self.assertNotEqual( batch_fingerprint.compute_fingerprint(a), batch_fingerprint.compute_fingerprint(b), ) def test_missing_data_returns_none(self): """没有 parquet 就算不出指纹 → None(调用方按新批次放行)。""" empty = self.root / "empty" empty.mkdir() self.assertIsNone(batch_fingerprint.compute_fingerprint(empty)) def test_record_and_lookup_roundtrip(self): fingerprint = "v:test=abc;d:def" batch_fingerprint.record_fingerprint("UMIGripper_AI", fingerprint, "ep_000003") self.assertEqual( batch_fingerprint.lookup_duplicate( "UMIGripper_AI", fingerprint, {"ep_000003"}), "ep_000003", ) def test_lookup_ignores_deleted_episode(self): """索引指向的 episode 已删除 → 不判重(允许重新上传)。""" fingerprint = "v:test=abc;d:def" batch_fingerprint.record_fingerprint("UMIGripper_AI", fingerprint, "ep_000003") self.assertIsNone( batch_fingerprint.lookup_duplicate( "UMIGripper_AI", fingerprint, {"ep_000004"}), ) def test_lookup_none_fingerprint_is_safe(self): self.assertIsNone( batch_fingerprint.lookup_duplicate("p", None, {"a"})) def test_record_drops_stale_entry_for_same_episode(self): """同名重传换了内容:旧指纹不该继续指向这个 episode。""" batch_fingerprint.record_fingerprint("p", "old-content", "ep_1") batch_fingerprint.record_fingerprint("p", "new-content", "ep_1") self.assertIsNone( batch_fingerprint.lookup_duplicate("p", "old-content", {"ep_1"})) self.assertEqual( batch_fingerprint.lookup_duplicate("p", "new-content", {"ep_1"}), "ep_1", ) def test_forget_fingerprint(self): batch_fingerprint.record_fingerprint("p", "abc", "ep_1") batch_fingerprint.forget_fingerprint("ep_1") self.assertIsNone( batch_fingerprint.lookup_duplicate("p", "abc", {"ep_1"})) def test_corrupt_index_is_tolerated(self): batch_fingerprint.INDEX_PATH.write_text("{not json", encoding="utf-8") self.assertIsNone( batch_fingerprint.lookup_duplicate("p", "abc", {"ep_1"})) batch_fingerprint.record_fingerprint("p", "abc", "ep_1") self.assertEqual( batch_fingerprint.lookup_duplicate("p", "abc", {"ep_1"}), "ep_1") def test_index_is_valid_json_with_episode_map(self): batch_fingerprint.record_fingerprint("proj", "fp1", "ep_1") data = json.loads(batch_fingerprint.INDEX_PATH.read_text(encoding="utf-8")) self.assertEqual(data["proj"]["fp1"], "ep_1") class ParquetHashTests(unittest.TestCase): def test_episode_index_column_is_excluded(self): with tempfile.TemporaryDirectory() as tmp: root = Path(tmp) table = pa.table({ "episode_index": pa.array([7], pa.int64()), "frame_index": pa.array([0], pa.int64()), "v": pa.array([[1.0, 2.0]], pa.list_(pa.float32())), }) one = root / "one.parquet" pq.write_table(table, one) other = root / "other.parquet" pq.write_table( table.set_column(0, "episode_index", pa.array([99], pa.int64())), other) self.assertEqual( batch_fingerprint._parquet_content_hash(one), batch_fingerprint._parquet_content_hash(other), ) def test_unreadable_parquet_returns_none(self): with tempfile.TemporaryDirectory() as tmp: bad = Path(tmp) / "bad.parquet" bad.write_bytes(b"not a parquet file") self.assertIsNone(batch_fingerprint._parquet_content_hash(bad)) if __name__ == "__main__": unittest.main()