test with real data

pylhc · Aug 20, 2024 · f571a3c · f571a3c
1 parent b3848b2
commit f571a3c
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 1 deletion.
diff --git a/tests/inputs/test_doros.h5 b/tests/inputs/test_doros.h5
diff --git a/tests/test_doros.py b/tests/test_doros.py
@@ -17,6 +17,23 @@
 INPUTS_DIR = Path(__file__).parent / "inputs"
 
 
+def test_read_write_real_data(tmp_path):
+    tbt = read_tbt(INPUTS_DIR / "test_doros.h5", bunch_id=10)
+
+    assert tbt.nbunches == 1
+    assert len(tbt.matrices) == 1
+    assert tbt.nturns == 50000
+    assert tbt.matrices[0].X.shape == (3, tbt.nturns)
+    assert tbt.matrices[0].Y.shape == (3, tbt.nturns)
+    assert len(set(tbt.matrices[0].X.index)) == 3
+    assert np.all(tbt.matrices[0].X.index == tbt.matrices[0].Y.index)
+
+    file_path = tmp_path / "test_file.h5"
+    write_tbt(tbt, file_path)
+    new = read_tbt(file_path, bunch_id=10)
+    compare_tbt(tbt, new, no_binary=False)
+
+
 def test_write_read(tmp_path):
     tbt = _tbt_data()
     file_path = tmp_path / "test_file.h5"

diff --git a/turn_by_turn/io.py b/turn_by_turn/io.py
@@ -28,7 +28,7 @@
     trackone=trackone,
     ascii=ascii,
 )
-WRITERS = ("lhc", "sps", "ascii")  # implemented writers
+WRITERS = ("lhc", "sps", "doros", "ascii")  # implemented writers
 
 write_lhc_ascii = write_ascii  # Backwards compatibility <0.4