"""trajectory npz 打包/解包单元测试（照搬旧项目 test_ceus_trajectory.py）。 旧标题：B11 trajectory npz 打包/解包单元测试。 钉死轨迹工件的单一 schema：完整往返还原全部键（沿用参考 ``write_roi_trajectory_npz`` 键集 + 新增 ``keyframe_indices`` / ``dark_polygons`` / ``calibration_json``）、 各键 dtype 精确（int32 / float32 / uint8 / str）、变长 ``dark_polygons``（逐帧 不同数量多边形）与 ``keyframe_indices`` 往返、``reference_region`` / ``calibration`` dict 往返，以及防御：坏字节解包 raise、形状不一致 pack raise、禁 pickle （含 object pickle 的恶意 npz 解包被拒）。测试用真实数据的二次构造（合成多边形 / 掩膜），边缘场景靠实现里的 assertion/raise 保证。 """ import io import pickle import numpy as np import pytest from cadx.ceus import ( Trajectory, ViewRect, pack_trajectory, unpack_trajectory, ) # Phase 1: 构造一个真实形态的 Trajectory（二次构造的合成几何） def _make_polygon(offset: float) -> np.ndarray: """构造 (64, 2) float32 闭合多边形（单位圆平移），与重采样输出形态一致。""" angles = np.linspace(0.0, 2.0 * np.pi, 64, endpoint=False) points = np.stack([np.cos(angles), np.sin(angles)], axis=1) * 10.0 + offset return points.astype(np.float32) def _make_trajectory() -> Trajectory: num_frames = 3 target_polygons = [_make_polygon(float(i)) for i in range(num_frames)] source_polygons = [_make_polygon(float(i) + 100.0) for i in range(num_frames)] # 逐帧暗区多边形数量不同：0 个 / 1 个 / 2 个（变长嵌套） dark_polygons: list[list[np.ndarray]] = [ [], [_make_polygon(5.0)], [_make_polygon(6.0), _make_polygon(7.0)], ] reference_mask = np.zeros((20, 30), dtype=np.uint8) reference_mask[5:10, 5:15] = 1 return Trajectory( frame_indices=np.arange(num_frames, dtype=np.int32), fps=29.97, target_polygons=np.asarray(target_polygons, dtype=np.float32), source_polygons=np.asarray(source_polygons, dtype=np.float32), source_view_rect=ViewRect(0, 0, 100, 200), target_view_rect=ViewRect(100, 0, 200, 200), source_effective_rect=ViewRect(5, 5, 95, 195), target_effective_rect=ViewRect(105, 5, 195, 195), source_polygon_coordinate_space="source_view_local", target_polygon_coordinate_space="target_effective_view_local", reference_mask=reference_mask, reference_region={"frame_index": 1, "area": 50.0, "source": "auto"}, keyframe_indices=np.asarray([1, 2], dtype=np.int32), dark_polygons=dark_polygons, calibration={ "source": "config/ceus.yaml", "config_hash": "abc123", "cli_override": False, "detected_split_x": 100, }, ) # Phase 2: 完整往返 —— 全部键还原 def test_roundtrip_preserves_all_keys() -> None: trajectory = _make_trajectory() restored = unpack_trajectory(pack_trajectory(trajectory)) np.testing.assert_array_equal(restored.frame_indices, trajectory.frame_indices) assert restored.fps == pytest.approx(trajectory.fps) np.testing.assert_allclose(restored.target_polygons, trajectory.target_polygons) np.testing.assert_allclose(restored.source_polygons, trajectory.source_polygons) assert restored.source_view_rect == trajectory.source_view_rect assert restored.target_view_rect == trajectory.target_view_rect assert restored.source_effective_rect == trajectory.source_effective_rect assert restored.target_effective_rect == trajectory.target_effective_rect assert restored.source_polygon_coordinate_space == trajectory.source_polygon_coordinate_space assert restored.target_polygon_coordinate_space == trajectory.target_polygon_coordinate_space np.testing.assert_array_equal(restored.reference_mask, trajectory.reference_mask) assert restored.reference_region == trajectory.reference_region np.testing.assert_array_equal(restored.keyframe_indices, trajectory.keyframe_indices) assert restored.calibration == trajectory.calibration # Phase 3: dtype 精确 def test_roundtrip_dtypes_are_exact() -> None: restored = unpack_trajectory(pack_trajectory(_make_trajectory())) assert restored.frame_indices.dtype == np.int32 assert restored.target_polygons.dtype == np.float32 assert restored.source_polygons.dtype == np.float32 assert restored.reference_mask.dtype == np.uint8 assert restored.keyframe_indices.dtype == np.int32 assert isinstance(restored.fps, float) assert isinstance(restored.source_polygon_coordinate_space, str) assert isinstance(restored.target_polygon_coordinate_space, str) def test_polygon_shapes_are_n_64_2() -> None: restored = unpack_trajectory(pack_trajectory(_make_trajectory())) assert restored.target_polygons.shape == (3, 64, 2) assert restored.source_polygons.shape == (3, 64, 2) # Phase 3b: 存档层 schema —— 直接在 npz 原始存档上断言键名 / dtype / shape # （对象层往返不足以守住 pack 的存储 schema；pack 若 schema 漂移须在此被发现） def test_archive_schema_keys_dtypes_shapes() -> None: """打开 pack 产出的原始 npz 存档，逐一断言键集 / 存储 dtype / shape。 与对象层往返互补：往返只验 unpack(pack(x))==x，无法发现 pack 写入存档时的 dtype / shape 漂移（如 fps 被写成 float64、polygons 退化）。此处用 ``allow_pickle=False`` 直接读取存档，钉死存储 schema。 """ archive_bytes = pack_trajectory(_make_trajectory()) with np.load(io.BytesIO(archive_bytes), allow_pickle=False) as archive: # 键集合 == 全部必要键（无多无少） expected_keys = { "frame_indices", "fps", "target_polygons", "source_polygons", "source_view_rect", "target_view_rect", "source_effective_rect", "target_effective_rect", "source_polygon_coordinate_space", "target_polygon_coordinate_space", "reference_mask", "reference_region_json", "keyframe_indices", "dark_polygons", "calibration_json", } assert set(archive.files) == expected_keys # 存档内 dtype 精确（存储 dtype 是 schema 的一部分） assert archive["frame_indices"].dtype == np.int32 assert archive["keyframe_indices"].dtype == np.int32 assert archive["target_polygons"].dtype == np.float32 assert archive["source_polygons"].dtype == np.float32 # fps 存档 dtype 必须 float32（即使 in-memory unpack 还原为 Python float） assert archive["fps"].dtype == np.float32 for rect_key in ( "source_view_rect", "target_view_rect", "source_effective_rect", "target_effective_rect", ): assert archive[rect_key].dtype == np.int32 assert archive["reference_mask"].dtype == np.uint8 # shape 精确 assert archive["target_polygons"].shape == (3, 64, 2) assert archive["source_polygons"].shape == (3, 64, 2) assert archive["frame_indices"].shape == (3,) for rect_key in ( "source_view_rect", "target_view_rect", "source_effective_rect", "target_effective_rect", ): assert archive[rect_key].shape == (4,) assert archive["reference_mask"].ndim == 2 # Phase 4: 变长 dark_polygons 往返 def test_roundtrip_variable_length_dark_polygons() -> None: trajectory = _make_trajectory() restored = unpack_trajectory(pack_trajectory(trajectory)) assert len(restored.dark_polygons) == 3 assert len(restored.dark_polygons[0]) == 0 assert len(restored.dark_polygons[1]) == 1 assert len(restored.dark_polygons[2]) == 2 for original_frame, restored_frame in zip( trajectory.dark_polygons, restored.dark_polygons, strict=True ): for original_poly, restored_poly in zip(original_frame, restored_frame, strict=True): assert isinstance(restored_poly, np.ndarray) assert restored_poly.dtype == np.float32 np.testing.assert_allclose(restored_poly, original_poly) def test_roundtrip_empty_dark_polygons() -> None: trajectory = _make_trajectory() trajectory.dark_polygons = [[], [], []] restored = unpack_trajectory(pack_trajectory(trajectory)) assert restored.dark_polygons == [[], [], []] # Phase 5: 防御 —— 坏字节 / 形状不一致 / 禁 pickle def test_unpack_bad_bytes_raises() -> None: with pytest.raises((ValueError, OSError)): unpack_trajectory(b"this-is-not-a-valid-npz-archive") def test_pack_rejects_mismatched_frame_count() -> None: trajectory = _make_trajectory() # frame_indices 长度与多边形帧数不一致 trajectory.frame_indices = np.arange(5, dtype=np.int32) with pytest.raises(ValueError): pack_trajectory(trajectory) def test_pack_rejects_wrong_polygon_shape() -> None: trajectory = _make_trajectory() # 把 target_polygons 改成 (N, 32, 2)，违反 (N, 64, 2) trajectory.target_polygons = np.zeros((3, 32, 2), dtype=np.float32) with pytest.raises(ValueError): pack_trajectory(trajectory) def test_pack_rejects_non_2d_reference_mask() -> None: trajectory = _make_trajectory() trajectory.reference_mask = np.zeros((20, 30, 3), dtype=np.uint8) with pytest.raises(ValueError): pack_trajectory(trajectory) def test_unpack_rejects_pickled_object_payload() -> None: """构造真实含 object pickle 的恶意 .npz（须 allow_pickle 才能 save/load），解包须被拒。 关键：用 ``np.savez_compressed`` 生成真正的 .npz 存档（而非 .npy），其中含一个 object dtype 数组（背后是 pickle payload）。这覆盖了「合法 .npz 容器、但成员需 pickle 反序列化」的注入面——``unpack_trajectory`` 内部 ``allow_pickle=False`` 必须 在访问该成员时 raise，绝不静默反序列化执行。 """ buffer = io.BytesIO() # 把 object dtype 数组放进 unpack 首先访问的必要键 frame_indices，确保命中 # allow_pickle=False 对 object 成员的拒绝路径（而非仅缺键 KeyError）。 # object dtype 数组在 savez 时即以 pickle 存储，构造真实 .npz 字节。 np.savez_compressed(buffer, frame_indices=np.array([{"evil": 1}], dtype=object)) raw = buffer.getvalue() # 访问 frame_indices 这一 object 成员时，allow_pickle=False 必须 raise（ValueError） with pytest.raises((ValueError, OSError, KeyError)): unpack_trajectory(raw) def test_pickle_payload_cannot_smuggle_through_unpack() -> None: """即便 pickle 字节伪装，unpack 内部 allow_pickle=False 也不得执行反序列化。""" payload = pickle.dumps({"evil": "code"}) with pytest.raises((ValueError, OSError, KeyError)): unpack_trajectory(payload)