Extend tests for split_trains slightly

European-XFEL · Oct 27, 2023 · 885c835 · 885c835
1 parent 7a9ddcb
commit 885c835
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/extra_data/tests/test_reader_mockdata.py b/extra_data/tests/test_reader_mockdata.py
@@ -702,6 +702,8 @@ def test_split_trains(mock_fxe_raw_run):
     chunks = list(run.split_trains(3))
     assert len(chunks) == 3
     assert {len(c.train_ids) for c in chunks} == {160}
+    arr = chunks[0]['FXE_XAD_GEC/CAM/CAMERA:daqOutput', 'data.image.dims'].ndarray()
+    assert arr.shape == (160, 2)
 
     chunks = list(run.split_trains(4, trains_per_part=100))
     assert len(chunks) == 5

diff --git a/extra_data/tests/test_sourcedata.py b/extra_data/tests/test_sourcedata.py
@@ -111,6 +111,8 @@ def test_split_trains(mock_spb_raw_run):
     chunks = list(xgm.split_trains(3))
     assert len(chunks) == 3
     assert {len(c.train_ids) for c in chunks} == {21, 22}
+    # The middle chunk spans across 2 files
+    assert [len(c.files) for c in chunks] == [1, 2, 1]
 
     chunks = list(xgm.split_trains(3, trains_per_part=20))
     assert len(chunks) == 4