implement simple test of DataSource.get_contiguous_time_periods() #223

Author: JackKelly

tests/data_sources/test_data_source.py

@@ -2,7 +2,6 @@
 
 
 def test_image_data_source():
-
     _ = ImageDataSource(
         image_size_pixels=64,
         meters_per_pixel=2000,

tests/data_sources/test_nwp_data_source.py

@@ -1,16 +1,17 @@
 import os
+import pandas as pd
 
 import nowcasting_dataset
 from nowcasting_dataset.data_sources.nwp.nwp_data_source import NWPDataSource
 
 
-def test_nwp_data_source_init():
+PATH = os.path.dirname(nowcasting_dataset.__file__)
 
-    path = os.path.dirname(nowcasting_dataset.__file__)
+# Solar PV data (test data)
+NWP_FILENAME = f"{PATH}/../tests/data/nwp_data/test.zarr"
 
-    # Solar PV data (test data)
-    NWP_FILENAME = f"{path}/../tests/data/nwp_data/test.zarr"
 
+def test_nwp_data_source_init():
     _ = NWPDataSource(
         filename=NWP_FILENAME,
         history_minutes=30,
@@ -21,12 +22,6 @@ def test_nwp_data_source_init():
 
 
 def test_nwp_data_source_open():
-
-    path = os.path.dirname(nowcasting_dataset.__file__)
-
-    # Solar PV data (test data)
-    NWP_FILENAME = f"{path}/../tests/data/nwp_data/test.zarr"
-
     nwp = NWPDataSource(
         filename=NWP_FILENAME,
         history_minutes=30,
@@ -40,12 +35,6 @@ def test_nwp_data_source_open():
 
 
 def test_nwp_data_source_batch():
-
-    path = os.path.dirname(nowcasting_dataset.__file__)
-
-    # Solar PV data (test data)
-    NWP_FILENAME = f"{path}/../tests/data/nwp_data/test.zarr"
-
     nwp = NWPDataSource(
         filename=NWP_FILENAME,
         history_minutes=30,
@@ -64,3 +53,20 @@ def test_nwp_data_source_batch():
     batch = nwp.get_batch(t0_datetimes=t0_datetimes, x_locations=x, y_locations=y)
 
     assert batch.data.shape == (4, 1, 19, 2, 2)
+
+
+def test_nwp_get_contiguous_time_periods():
+    nwp = NWPDataSource(
+        filename=NWP_FILENAME,
+        history_minutes=30,
+        forecast_minutes=60,
+        convert_to_numpy=True,
+        n_timesteps_per_batch=8,
+        channels=["t"],
+    )
+
+    contiguous_time_periods = nwp.get_contiguous_time_periods()
+    correct_time_periods = pd.DataFrame(
+        [{"start_dt": pd.Timestamp("2019-01-01 00:00"), "end_dt": pd.Timestamp("2019-01-02 02:00")}]
+    )
+    pd.testing.assert_frame_equal(contiguous_time_periods, correct_time_periods)