"""Persistence invariants: complete data, compatible sets and recoverable write failures.""" from __future__ import annotations import json from pathlib import Path import tempfile import unittest from unittest import mock import numpy as np from python_app.models.dataset_model import ( ComboKey, ResultBlock, ResultCollection, ResultPayload, SweepCollection, TraceData, ) from python_app.storage.npz.store import NpzStore from python_app.storage.npz.snapshot_numpy import select_aligned_histories from python_app.storage.npz.vna_history_json import build_vna_history_payload from python_app.scripts.replay_raw_producer import _load_dataset from python_app.scripts.convert_snapshot_to_vna_history import _collection_dirs, _load_stage_records, _build_sweep_history from python_app.workflows.sequential_capture_workflow import combine_collections_via_median, combine_traces_via_median def _sweep(cid: int = 1, ns: int = 20, *, start: int = 10, end: int = 15) -> SweepCollection: return SweepCollection( collection_id=cid, monotonic_ns=ns, capture_start_ns=start, capture_end_ns=end, traces=[TraceData( combo=ComboKey(0, 0), frequency_hz=np.array([1e9, 2e9], dtype=np.float32), s11=np.array([1 + 2j, 3 - 4j], dtype=np.complex64), s21=np.array([5 - 6j, 7 + 8j], dtype=np.complex64), capture_start_ns=start, capture_end_ns=end, )], ) def _payloads() -> list[ResultPayload]: return [ ResultPayload('trace', 1, frequency_hz=np.array([1, 2], dtype=np.float32), trace=np.array([1 + 2j, 3 - 4j], dtype=np.complex64)), ResultPayload('scalar', 2, scalar_value=3.5), ResultPayload('image', 3, image_x_axis=np.array([1, 2], dtype=np.float32), image_y_axis=np.array([3], dtype=np.float32), image=np.array([[4, 5]], dtype=np.float32)), ResultPayload('table', 4, table=np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32)), ] class PersistenceIntegrityTest(unittest.TestCase): def setUp(self) -> None: temporary = tempfile.TemporaryDirectory() self.addCleanup(temporary.cleanup) self.root = Path(temporary.name) self.store = NpzStore(self.root / 'sets') self.output = self.root / 'output' def test_all_calibration_and_reference_kinds_preserve_arrays_and_timing(self) -> None: for kind in ('s21_calibration', 's21_reference', 's11_open', 's11_short', 's11_load', 's11_reference'): with self.subTest(kind=kind): original = _sweep(7, 100) second = _sweep().traces[0] second.combo = ComboKey(2, 3) original.traces.append(second) self.store.save_set(kind, 'radar', 'full', original) loaded = self.store.load_set(kind, 'radar', 'full') self.assertEqual((loaded.collection_id, loaded.monotonic_ns), (7, 100)) self.assertEqual((loaded.capture_start_ns, loaded.capture_end_ns), (10, 15)) self.assertTrue(self.store.has_combo_coverage(kind, 'radar', 'full', [ComboKey(0, 0), ComboKey(2, 3)])) self.assertFalse(self.store.has_combo_coverage(kind, 'radar', 'full', [ComboKey(9, 9)])) for before, after in zip(original.traces, loaded.traces, strict=True): self.assertEqual(before.combo, after.combo) self.assertEqual((after.capture_start_ns, after.capture_end_ns), (10, 15)) for field in ('frequency_hz', 's11', 's21'): np.testing.assert_array_equal(getattr(before, field), getattr(after, field)) def test_new_set_remains_loadable_without_json_sidecar(self) -> None: self.store.save_set('s21_reference', 'radar', 'full', _sweep()) (self.root / 'sets/s21_reference/radar/full.json').unlink() loaded = self.store.load_set('s21_reference', 'radar', 'full') np.testing.assert_array_equal(loaded.traces[0].s21, _sweep().traces[0].s21) def test_legacy_npz_json_pair_is_still_loadable(self) -> None: self.store.save_set('s21_reference', 'radar', 'full', _sweep()) archive = self.root / 'sets/s21_reference/radar/full.npz' with np.load(archive, allow_pickle=False) as arrays: legacy = {key: arrays[key] for key in arrays.files if key != 'metadata_json'} np.savez(archive, **legacy) self.assertEqual(self.store.load_set('s21_reference', 'radar', 'full').monotonic_ns, 20) def test_failed_set_commit_preserves_previous_data_and_sidecar(self) -> None: self.store.save_set('s21_reference', 'radar', 'full', _sweep(1, 20)) original_meta = (self.root / 'sets/s21_reference/radar/full.json').read_bytes() replace = Path.replace def fail_archive_commit(path, destination): if Path(destination).suffix == '.npz': raise OSError('disk full during commit') return replace(path, destination) with mock.patch.object(Path, 'replace', fail_archive_commit), self.assertRaises(OSError): self.store.save_set('s21_reference', 'radar', 'full', _sweep(2, 200)) self.assertEqual(self.store.load_set('s21_reference', 'radar', 'full').collection_id, 1) self.assertEqual((self.root / 'sets/s21_reference/radar/full.json').read_bytes(), original_meta) self.assertEqual(list(self.root.rglob('*.tmp')), []) def test_failed_new_set_commit_does_not_publish_a_sidecar_only_set(self) -> None: replace = Path.replace def fail_archive_commit(path, destination): if Path(destination).suffix == '.npz': raise OSError('commit failed') return replace(path, destination) with mock.patch.object(Path, 'replace', fail_archive_commit), self.assertRaises(OSError): self.store.save_set('s21_reference', 'radar', 'full', _sweep()) self.assertEqual(self.store.list_sets('s21_reference', 'radar'), []) self.assertEqual(list(self.root.rglob('*.json')), []) self.assertEqual(list(self.root.rglob('*.tmp')), []) def test_failed_sidecar_commit_restores_previous_archive_without_rewriting_it(self) -> None: self.store.save_set('s21_reference', 'radar', 'full', _sweep(1, 20)) original = (self.root / 'sets/s21_reference/radar/full.npz').read_bytes() replace = Path.replace def fail_metadata_commit(path, destination): if Path(destination).suffix == '.json': raise OSError('metadata commit failed') return replace(path, destination) with mock.patch.object(Path, 'replace', fail_metadata_commit), self.assertRaises(OSError): self.store.save_set('s21_reference', 'radar', 'full', _sweep(2, 200)) self.assertEqual((self.root / 'sets/s21_reference/radar/full.npz').read_bytes(), original) self.assertEqual(self.store.load_set('s21_reference', 'radar', 'full').collection_id, 1) self.assertEqual(list(self.root.rglob('*.tmp')), []) def test_invalid_sets_are_rejected_before_overwriting_valid_data(self) -> None: self.store.save_set('s21_calibration', 'radar', 'full', _sweep()) invalid = [_sweep(), _sweep(), _sweep(), _sweep()] invalid[0].traces.append(invalid[0].traces[0]) invalid[1].traces[0].s21 = np.array([1], dtype=np.complex64) invalid[2].traces[0].frequency_hz = np.array([[1, 2]], dtype=np.float32) invalid[3].traces[0].s11[0] = complex(float('nan'), 0) for collection in invalid: with self.subTest(collection=collection), self.assertRaises(ValueError): self.store.save_set('s21_calibration', 'radar', 'full', collection) self.assertEqual(self.store.load_set('s21_calibration', 'radar', 'full').monotonic_ns, 20) def test_set_names_cannot_escape_store(self) -> None: for name in ('../escape', 'sub/name', '..', '.', '', r'sub\name'): with self.subTest(name=name), self.assertRaises(ValueError): self.store.save_set('s21_reference', 'radar', name, _sweep()) def test_invalid_channel_identifiers_are_rejected_without_truncation(self) -> None: for position in (0.5, True, -1, 2 ** 32): sweep = _sweep() sweep.traces[0].combo = ComboKey(position, 0) with self.subTest(position=position), self.assertRaises(ValueError): self.store.save_set('s21_reference', 'radar', 'invalid', sweep) self.assertEqual(self.store.list_sets('s21_reference', 'radar'), []) def test_corrupted_metadata_types_are_rejected_without_truncation(self) -> None: self.store.save_set('s21_reference', 'radar', 'full', _sweep()) archive = self.root / 'sets/s21_reference/radar/full.npz' with np.load(archive, allow_pickle=False) as stored: arrays = {key: stored[key] for key in stored.files} original = json.loads(str(arrays['metadata_json'].item())) for malformed in ([], {**original, 'collection_id': 1.5}, {**original, 'monotonic_ns': True}, {**original, 'combos': [{**original['combos'][0], 'input': 0.5}]}): with self.subTest(metadata=malformed): np.savez(archive, **{**arrays, 'metadata_json': np.asarray(json.dumps(malformed))}) with self.assertRaisesRegex(RuntimeError, 'Corrupted preprocess set'): self.store.load_set('s21_reference', 'radar', 'full') def test_complete_snapshot_keeps_unmatched_stage_tails(self) -> None: raw = [_sweep(1, 10), _sweep(2, 30)] pre = [_sweep(1, 20), _sweep(3, 40)] result = [ResultCollection(1, 20)] directory, summary = self.store.save_runtime_snapshot_numpy( self.output, 'full', raw, pre, result, 1000, include_unaligned=True, ) self.assertEqual((summary['raw_count'], summary['preprocessed_count'], summary['result_count']), (2, 2, 1)) for stage, ids in [('raw', [1, 2]), ('preprocessed', [1, 3]), ('results', [1])]: metadata = [json.loads(path.read_text()) for path in sorted((directory / stage).glob('*/meta.json'))] self.assertEqual([item['collection_id'] for item in metadata], ids) def test_snapshot_result_payloads_preserve_every_kind_at_both_scopes(self) -> None: payloads = _payloads() result = ResultCollection(1, 20, processing_duration_ns=12345, collection_payloads=payloads, blocks=[ResultBlock(ComboKey(0, 0), payloads)]) directory, _ = self.store.save_runtime_snapshot_numpy(self.output, 'full', [], [], [result], 1000) collection_dir = next((directory / 'results').iterdir()) meta = json.loads((collection_dir / 'meta.json').read_text()) self.assertEqual(meta['processing_duration_ns'], 12345) for base, entries in [(collection_dir, meta['collection_payloads']), (collection_dir / meta['blocks'][0]['dir'], meta['blocks'][0]['payloads'])]: self.assertEqual([entry['kind'] for entry in entries], [1, 2, 3, 4]) np.testing.assert_array_equal(np.load(base / entries[0]['freq_file']), payloads[0].frequency_hz) np.testing.assert_array_equal(np.load(base / entries[0]['trace_file']), payloads[0].trace) np.testing.assert_array_equal(np.load(base / entries[1]['scalar_file']), np.array([3.5], dtype=np.float32)) np.testing.assert_array_equal(np.load(base / entries[2]['x_axis_file']), payloads[2].image_x_axis) np.testing.assert_array_equal(np.load(base / entries[2]['y_axis_file']), payloads[2].image_y_axis) np.testing.assert_array_equal(np.load(base / entries[2]['image_file']), payloads[2].image) np.testing.assert_array_equal(np.load(base / entries[3]['table_file']), payloads[3].table) def test_bad_result_payload_rolls_back_entire_snapshot(self) -> None: for payload in (ResultPayload('unknown', 9), ResultPayload('bad_table', 4, table=np.array([1, 2]))): with self.subTest(kind=payload.kind), self.assertRaises(ValueError): self.store.save_runtime_snapshot_numpy( self.output, 'failed', [_sweep()], [], [ResultCollection(1, 20, collection_payloads=[payload])], 1000, ) self.assertFalse((self.output / 'failed').exists()) self.assertEqual(list(self.output.iterdir()), []) def test_snapshot_write_failure_can_be_retried_with_same_name(self) -> None: with mock.patch('numpy.save', side_effect=OSError('disk full')), self.assertRaises(OSError): self.store.save_runtime_snapshot_numpy(self.output, 'retry', [_sweep()], [], [], 1000) directory, _ = self.store.save_runtime_snapshot_numpy(self.output, 'retry', [_sweep()], [], [], 1000) self.assertTrue(directory.exists()) with self.assertRaises(FileExistsError): self.store.save_runtime_snapshot_numpy(self.output, 'retry', [_sweep()], [], [], 1000) def test_stream_failure_keeps_prior_chunks_and_retries_without_partial_data(self) -> None: writer = self.store.create_snapshot_stream(self.output, 'stream') writer.append([_sweep()], [], [ResultCollection(1, 20)]) with mock.patch('numpy.save', side_effect=OSError('disk full')), self.assertRaises(OSError): writer.append([_sweep(2, 30)], [], [ResultCollection(2, 40)]) self.assertEqual(writer.result_count, 1) self.assertEqual(len(list((writer.directory / 'raw').iterdir())), 1) self.assertEqual(list(writer.directory.glob('.chunk-*')), []) writer.append([_sweep(2, 30)], [], [ResultCollection(2, 40)]) self.assertEqual(writer.result_count, 2) self.assertEqual([path.name.split('_')[0] for path in sorted((writer.directory / 'raw').iterdir())], ['0000', '0001']) def test_json_export_preserves_unmatched_raw_sweeps_and_matches_reused_ids_by_timing(self) -> None: raw_old = _sweep(1, 10, start=1, end=5) raw_new = _sweep(1, 100, start=90, end=95) raw_new.traces[0].s21[:] = 99 pre_new = _sweep(1, 110, start=90, end=95) payload = build_vna_history_payload([raw_old, raw_new], [pre_new], [], input_index=0, output_index=0) self.assertEqual(len(payload['sweep_history']), 2) self.assertEqual(payload['sweep_history'][-1]['sweep_points'], [[99.0, 0.0], [99.0, 0.0]]) self.assertEqual(payload['sweep_history'][-1]['calibrated_points'], [[5.0, -6.0], [7.0, 8.0]]) def test_json_batch_failure_does_not_publish_partial_export(self) -> None: sweep = _sweep() second = _sweep().traces[0] second.combo = ComboKey(0, 1) sweep.traces.append(second) write_text = Path.write_text writes = 0 def fail_second_file(path, *args, **kwargs): nonlocal writes writes += 1 if writes == 2: raise OSError('disk full') return write_text(path, *args, **kwargs) with mock.patch.object(Path, 'write_text', fail_second_file), self.assertRaises(OSError): self.store.save_runtime_vna_history_json_batch(self.output, 'export', [sweep], [], [], 1000) self.assertEqual(list(self.output.iterdir()), []) paths, summary = self.store.save_runtime_vna_history_json_batch(self.output, 'export', [sweep], [], [], 1000, channel='s11') self.assertEqual(summary['combo_count'], 2) self.assertEqual(len(paths), 2) for path in paths: self.assertEqual(json.loads(path.read_text())['channel'], 's11') def test_saved_dataset_replay_preserves_s11_and_accepts_legacy_s21_only(self) -> None: directory, _ = self.store.save_runtime_snapshot_numpy(self.output, 'replay', [_sweep()], [], [], 1000) dataset = _load_dataset(directory / 'raw', [ComboKey(0, 0)], 2) np.testing.assert_array_equal(dataset.snapshots[0][0].s11, _sweep().traces[0].s11) self.assertEqual(dataset.snapshots[0][0].capture_end_ns, 15) next((directory / 'raw').glob('*/i0_o0_s11.npy')).unlink() legacy = _load_dataset(directory / 'raw', [ComboKey(0, 0)], 2) np.testing.assert_array_equal(legacy.snapshots[0][0].s11, np.zeros(2, dtype=np.complex64)) def test_long_recording_replay_and_converter_use_numeric_directory_order(self) -> None: writer = self.store.create_snapshot_stream(self.output, 'long') writer.append([_sweep(1, 10), _sweep(2, 20), _sweep(3, 30)], [], []) folders = sorted((writer.directory / 'raw').iterdir()) for folder, index in zip(folders, (9998, 9999, 10000), strict=True): folder.rename(folder.with_name(f'{index}_{folder.name.split("_", 1)[1]}')) dataset = _load_dataset(writer.directory / 'raw', [ComboKey(0, 0)], 2) self.assertEqual(dataset.timestamps_ns, [10, 20, 30]) self.assertEqual([int(path.name.split('_')[0]) for path in _collection_dirs(writer.directory / 'raw')], [9998, 9999, 10000]) def test_offline_and_runtime_json_exports_share_matching_and_preserve_orphan_sweeps(self) -> None: raw = [_sweep(1, 10, start=1, end=5), _sweep(1, 100, start=90, end=95)] pre = [_sweep(1, 110, start=90, end=95)] directory, _ = self.store.save_runtime_snapshot_numpy( self.output, 'convert', raw, pre, [], 1000, include_unaligned=True, ) records = { stage: _load_stage_records(directory, stage, input_index=0, output_index=0, channel='s21') for stage in ('raw', 'preprocessed') } offline = _build_sweep_history(records['raw'], records['preprocessed'], channel='s21', primary_stage='preprocessed') runtime = build_vna_history_payload(raw, pre, [], input_index=0, output_index=0) self.assertEqual(offline, runtime['sweep_history']) self.assertEqual(len(offline), 2) class SnapshotAlignmentTest(unittest.TestCase): def test_reused_ids_and_missing_frames_do_not_cross_align_runs(self) -> None: raw_old = _sweep(1, 10, start=1, end=5) pre_old = _sweep(1, 20, start=1, end=5) pre_new = _sweep(1, 200, start=150, end=175) selected_raw, selected_pre, results, summary = select_aligned_histories( [raw_old], [pre_old, pre_new], [ResultCollection(1, 200)], 1, ) self.assertEqual(selected_raw, []) self.assertEqual(selected_pre, [pre_new]) self.assertEqual(summary['raw_missing_count'], 1) self.assertEqual(results[0].monotonic_ns, 200) def test_result_does_not_pair_with_wrong_preprocessed_timestamp(self) -> None: _, pre, _, summary = select_aligned_histories([], [_sweep(1, 20)], [ResultCollection(1, 200)], 1000) self.assertEqual(pre, []) self.assertEqual(summary['preprocessed_missing_count'], 1) def test_independent_stage_limit_is_exactly_requested_tail(self) -> None: raw = [_sweep(cid=index, ns=index) for index in range(1005)] selected_raw, _, _, _ = select_aligned_histories(raw, [], [], 1000, include_unaligned=True) self.assertEqual(len(selected_raw), 1000) self.assertEqual((selected_raw[0].collection_id, selected_raw[-1].collection_id), (5, 1004)) class CalibrationMedianIntegrityTest(unittest.TestCase): def test_rejects_frequency_axes_that_changed_between_captures(self) -> None: first, second = _sweep(), _sweep() second.traces[0].frequency_hz[1] += 1e8 with self.assertRaisesRegex(RuntimeError, 'frequency-axis'): combine_traces_via_median([first.traces[0], second.traces[0]]) def test_rejects_extra_combinations_instead_of_silently_discarding_them(self) -> None: first, second = _sweep(), _sweep() extra = _sweep().traces[0] extra.combo = ComboKey(0, 1) second.traces.append(extra) with self.assertRaisesRegex(RuntimeError, 'trace counts'): combine_collections_via_median([first, second]) def test_collection_and_trace_windows_span_all_median_captures(self) -> None: first, second = _sweep(start=10, end=15), _sweep(start=30, end=35) combined = combine_collections_via_median([first, second]) self.assertEqual((combined.capture_start_ns, combined.capture_end_ns), (10, 35)) self.assertEqual((combined.traces[0].capture_start_ns, combined.traces[0].capture_end_ns), (10, 35)) if __name__ == '__main__': unittest.main()