add tests

Author: peterdudfield

nowcasting_dataset/data_sources/sun/sun_data_source.py

@@ -127,14 +127,16 @@ def get_locations(
     def datetime_index(self) -> pd.DatetimeIndex:
         """Get datetimes where elevation >= 10"""
 
-        if not self.load_live:
-            # get the lat and lon from london
-            latitude = 51
-            longitude = 0
+        # get the lat and lon from london
+        latitude = 51
+        longitude = 0
 
+        if not self.load_live:
             datestamps = self.elevation.index
         else:
-            datestamps = pd.date_range(datetime(2019, 1, 1), datetime(2020, 1, 1), frew="5T")
+            datestamps = pd.date_range(
+                datetime(2019, 1, 1), datetime(2019, 12, 31, 23, 55), freq="5T"
+            )
 
         # get elevation for all datetimes
         azimuth_elevation = calculate_azimuth_and_elevation_angle(
@@ -152,7 +154,7 @@ def datetime_index(self) -> pd.DatetimeIndex:
             f"out of {len(azimuth_elevation)} as elevation is < 10"
         )
 
-        datetimes = self.elevation[mask].index
+        datetimes = datestamps[mask]
 
         # Sun data is only for 2019, so to expand on these by
         # repeating data from 2014 to 2023

tests/data_sources/sun/test_sun_data_source.py

@@ -8,7 +8,8 @@
 def test_init(test_data_folder):  # noqa 103
     zarr_path = test_data_folder + "/sun/test.zarr"
 
-    _ = SunDataSource(zarr_path=zarr_path, history_minutes=30, forecast_minutes=60)
+    sun = SunDataSource(zarr_path=zarr_path, history_minutes=30, forecast_minutes=60)
+    _ = sun.datetime_index()
 
 
 def test_get_example(test_data_folder):  # noqa 103
@@ -50,6 +51,7 @@ def test_get_load_live():  # noqa 103
     sun_data_source = SunDataSource(
         zarr_path="", history_minutes=30, forecast_minutes=60, load_live=True
     )
+    _ = sun_data_source.datetime_index()
 
     x = 256895.63164759654
     y = 666180.3018829626