diff --git a/protozfits/tests/test_R1_example_file.py b/protozfits/tests/test_R1_example_file.py index f6cd53f..778a314 100644 --- a/protozfits/tests/test_R1_example_file.py +++ b/protozfits/tests/test_R1_example_file.py @@ -1,8 +1,9 @@ - import numpy as np import pkg_resources import os +from protozfits import SimpleFile + example_file_path = pkg_resources.resource_filename( 'protozfits', os.path.join( @@ -13,20 +14,16 @@ ) -def test_can_open_file(): - from protozfits import SimpleFile - SimpleFile(example_file_path) - - def test_can_iterate_over_events_and_run_header(): - from protozfits import SimpleFile - f = SimpleFile(example_file_path) - camera_config = next(f.CameraConfig) - assert (camera_config.expected_pixels_id == np.arange(14)).all() + with SimpleFile(example_file_path) as f: + + camera_config = next(f.CameraConfig) + assert (camera_config.expected_pixels_id == np.arange(14)).all() - for event in f.Events: - assert event.waveform.shape == (1120,) - assert event.pixel_status.shape == (14,) - assert event.lstcam.first_capacitor_id.shape == (16,) - assert event.lstcam.counters.shape == (44,) + for i, event in enumerate(f.Events): + assert event.event_id == i + 1 + assert event.waveform.shape == (1120,) + assert event.pixel_status.shape == (14,) + assert event.lstcam.first_capacitor_id.shape == (16,) + assert event.lstcam.counters.shape == (44,)