tests/unit_tests/test_dex.py

0001 # test_dex.py - DEX v1 upgrade, validation, and the upgrade CLI command
0002
0003 import json
0004 import os
0005 import zipfile
0006
0007 import pytest
0008 from click.testing import CliRunner
0009
0010 from pyrobird.cli.upgrade import upgrade
0011 from pyrobird.dex import (
0012     collect_dex_problems,
0013     make_dex,
0014     upgrade_dex,
0015     validate_dex,
0016     UnknownPieceTypeError,
0017 )
0018
0019 # A small but complete 0.04 document covering both known group types
0020 SAMPLE_004 = {
0021     "type": "firebird-dex-json",
0022     "version": "0.04",
0023     "origin": {"file": "sample.root"},
0024     "events": [
0025         {
0026             "id": "event_0",
0027             "groups": [
0028                 {
0029                     "name": "BarrelHits",
0030                     "type": "BoxHit",
0031                     "origin": {"type": "edm4eic::TrackerHitData"},
0032                     "hits": [
0033                         {"pos": [1, 2, 3], "dim": [0.1, 0.2, 0.3], "t": [4, 0.5], "ed": [0.001, 0.0001]},
0034                         {"pos": [7, 8, 9], "dim": [1, 1, 1], "t": [5, 0], "ed": [0.002, 0]},
0035                     ],
0036                 },
0037                 {
0038                     "name": "CentralTracks",
0039                     "type": "PointTrajectory",
0040                     "paramColumns": ["theta", "phi"],
0041                     "pointColumns": ["x", "y", "z", "t"],
0042                     "trajectories": [
0043                         {"points": [[0, 0, 0, 0], [1, 1, 1, 1]], "params": [0.5, 1.5]},
0044                         {"points": [[2, 2, 2, 2]], "params": [0.7, 2.5]},
0045                     ],
0046                 },
0047             ],
0048         }
0049     ],
0050 }
0051
0052
0053 def test_upgrade_box_hits():
0054     result = upgrade_dex(SAMPLE_004)
0055     validate_dex(result)
0056
0057     assert result["version"] == "1.0"
0058     event = result["events"][0]
0059     assert event["id"] == "event_0"
0060
0061     hits = next(p for p in event["pieces"] if p["name"] == "BarrelHits")
0062     assert hits["type"] == "BoxHit"
0063     assert hits["count"] == 2
0064     # hit 0 and hit 1 keep their order: id == index
0065     assert hits["columns"]["pos"] == [1, 2, 3, 7, 8, 9]
0066     assert hits["columns"]["dim"] == [0.1, 0.2, 0.3, 1, 1, 1]
0067     assert hits["columns"]["time"] == [4, 5]
0068     assert hits["columns"]["timeError"] == [0.5, 0]
0069     assert hits["columns"]["edep"] == [0.001, 0.002]
0070     assert hits["columns"]["edepError"] == [0.0001, 0]
0071     assert hits["origin"] == {"type": "edm4eic::TrackerHitData"}
0072
0073
0074 def test_upgrade_trajectories():
0075     result = upgrade_dex(SAMPLE_004)
0076     tracks = next(p for p in result["events"][0]["pieces"] if p["name"] == "CentralTracks")
0077
0078     assert tracks["type"] == "PointTrajectory"
0079     assert tracks["count"] == 2
0080     # per-trajectory params became per-name columns, same index
0081     assert tracks["columns"]["theta"] == [0.5, 0.7]
0082     assert tracks["columns"]["phi"] == [1.5, 2.5]
0083     assert tracks["pointColumns"] == ["x", "y", "z", "t"]
0084     assert tracks["points"][0] == [[0, 0, 0, 0], [1, 1, 1, 1]]
0085     assert tracks["points"][1] == [[2, 2, 2, 2]]
0086
0087
0088 def test_upgrade_old_tracker_line_type_name():
0089     # The pre-0.04 type name and its "lines" container map to PointTrajectory
0090     old = {
0091         "version": "0.03",
0092         "events": [{
0093             "id": 0,
0094             "groups": [{
0095                 "name": "Tracks",
0096                 "type": "TrackerLinePointTrajectory",
0097                 "paramColumns": [],
0098                 "pointColumns": ["x", "y", "z"],
0099                 "lines": [{"points": [[1, 2, 3]], "params": []}],
0100             }],
0101         }],
0102     }
0103     result = upgrade_dex(old)
0104     validate_dex(result)
0105     piece = result["events"][0]["pieces"][0]
0106     assert piece["type"] == "PointTrajectory"
0107     assert piece["count"] == 1
0108     assert piece["points"] == [[[1, 2, 3]]]
0109
0110
0111 def test_upgrade_unknown_type_fails_by_default():
0112     doc = {
0113         "version": "0.04",
0114         "events": [{"id": 0, "groups": [{"name": "X", "type": "custom.Unknown", "stuff": []}]}],
0115     }
0116     with pytest.raises(UnknownPieceTypeError):
0117         upgrade_dex(doc)
0118
0119     # skip_unknown drops the group and produces a valid (empty) event
0120     result = upgrade_dex(doc, skip_unknown=True)
0121     validate_dex(result)
0122     assert result["events"][0]["pieces"] == []
0123
0124
0125 def test_upgrade_is_a_noop_on_v1():
0126     v1 = make_dex([{"id": 0, "pieces": []}])
0127     assert upgrade_dex(v1) is v1
0128
0129
0130 def test_validate_catches_column_length_mismatch():
0131     doc = make_dex([{
0132         "id": 0,
0133         "pieces": [{
0134             "name": "Hits", "type": "BoxHit", "version": "1.0",
0135             "count": 2,
0136             "columns": {"pos": [1, 2, 3], "dim": [1, 1, 1, 2, 2, 2]},  # pos too short
0137         }],
0138     }])
0139     problems = collect_dex_problems(doc)
0140     assert any("pos" in p for p in problems)
0141     with pytest.raises(ValueError):
0142         validate_dex(doc)
0143
0144
0145 def test_validate_catches_bad_refs():
0146     doc = make_dex([{
0147         "id": 0,
0148         "pieces": [
0149             {
0150                 "name": "Particles", "type": "custom.Particle", "version": "1.0",
0151                 "count": 2,
0152                 "columns": {"pdg": [11, 22], "parent": [-1, 0]},
0153                 "refs": {"parent": "Particles"},
0154             },
0155             {
0156                 "name": "Broken", "type": "custom.Thing", "version": "1.0",
0157                 "count": 1,
0158                 "columns": {"particle": [5]},  # out of range for Particles (count=2)
0159                 "refs": {"particle": "Particles"},
0160             },
0161             {
0162                 "name": "AlsoBroken", "type": "custom.Thing2", "version": "1.0",
0163                 "count": 1,
0164                 "columns": {},
0165                 "refs": {"missing": "Nowhere"},  # no such column, no such piece
0166             },
0167         ],
0168     }])
0169     problems = collect_dex_problems(doc)
0170     assert any("not a valid index" in p for p in problems)
0171     assert any("no such column" in p for p in problems)
0172
0173     # the self-referencing Particles piece alone is valid
0174     solo = make_dex([{"id": 0, "pieces": [doc["events"][0]["pieces"][0]]}])
0175     validate_dex(solo)
0176
0177
0178 def test_validate_catches_ragged_points_mismatch():
0179     doc = make_dex([{
0180         "id": 0,
0181         "pieces": [{
0182             "name": "Tracks", "type": "PointTrajectory", "version": "1.0",
0183             "count": 2,
0184             "columns": {},
0185             "pointColumns": ["x", "y", "z"],
0186             "points": [[[1, 2, 3]]],  # only one entity, count says 2
0187         }],
0188     }])
0189     problems = collect_dex_problems(doc)
0190     assert any("points" in p for p in problems)
0191
0192
0193 def test_upgrade_cli_json_roundtrip(tmp_path):
0194     input_file = tmp_path / "old.firebird.json"
0195     output_file = tmp_path / "new.firebird.json"
0196     input_file.write_text(json.dumps(SAMPLE_004))
0197
0198     runner = CliRunner()
0199     result = runner.invoke(upgrade, [str(input_file), str(output_file)])
0200     assert result.exit_code == 0, result.output
0201
0202     upgraded = json.loads(output_file.read_text())
0203     validate_dex(upgraded)
0204     assert upgraded["version"] == "1.0"
0205
0206
0207 def test_upgrade_cli_zip_roundtrip(tmp_path):
0208     input_file = tmp_path / "old.firebird.zip"
0209     output_file = tmp_path / "new.firebird.zip"
0210     with zipfile.ZipFile(input_file, "w", compression=zipfile.ZIP_DEFLATED) as zf:
0211         zf.writestr("old.firebird.json", json.dumps(SAMPLE_004))
0212
0213     runner = CliRunner()
0214     result = runner.invoke(upgrade, [str(input_file), str(output_file)])
0215     assert result.exit_code == 0, result.output
0216
0217     with zipfile.ZipFile(output_file) as zf:
0218         member = [n for n in zf.namelist() if n.endswith(".json")][0]
0219         upgraded = json.loads(zf.read(member))
0220     validate_dex(upgraded)
0221     assert upgraded["version"] == "1.0"
0222
0223
0224 def test_upgrade_cli_default_output_name(tmp_path):
0225     input_file = tmp_path / "sample.firebird.json"
0226     input_file.write_text(json.dumps(SAMPLE_004))
0227
0228     runner = CliRunner()
0229     result = runner.invoke(upgrade, [str(input_file)])
0230     assert result.exit_code == 0, result.output
0231
0232     expected = tmp_path / "sample.v1.firebird.json"
0233     assert expected.exists()
0234     validate_dex(json.loads(expected.read_text()))
0235
0236
0237 def test_upgrade_cli_unknown_type(tmp_path):
0238     doc = {
0239         "version": "0.04",
0240         "events": [{"id": 0, "groups": [{"name": "X", "type": "custom.Unknown"}]}],
0241     }
0242     input_file = tmp_path / "custom.firebird.json"
0243     input_file.write_text(json.dumps(doc))
0244
0245     runner = CliRunner()
0246     result = runner.invoke(upgrade, [str(input_file)])
0247     assert result.exit_code != 0
0248     assert "custom.Unknown" in result.output
0249
0250     result = runner.invoke(upgrade, ["--skip-unknown", str(input_file)])
0251     assert result.exit_code == 0, result.output