diff --git a/tests/test_tools.py b/tests/test_tools.py index 66bac913e49e624141addf99be3c5d6517657d33..8168a3bed4ec2b211ec548210b2959530de8257e 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -183,6 +183,16 @@ class TestRecStagesMasks(unittest.TestCase): assert masks[1][0] == all(rec_stages[1][0] == ak.Array(stages)) assert masks[0][1] == False + def test_mask_with_start_and_end_of_rec_stages(self): + rec_stages = self.tracks.rec_stages + stages = [1, 3, 5, 4] + masks = mask(rec_stages, start=1, end=4) + + assert masks[0][0] == all(rec_stages[0][0] == ak.Array(stages)) + assert masks[1][0] == all(rec_stages[1][0] == ak.Array(stages)) + assert masks[0][1] == False + + class TestUnique(unittest.TestCase): def run_random_test_with_dtype(self, dtype):