More test assertions

api/app/tests/sfms/test_bui_date_range_processor.py

@@ -2,18 +2,20 @@
 from typing import List
 from unittest.mock import AsyncMock
 import pytest
-from datetime import datetime, timezone
+from datetime import datetime, timezone, timedelta
 from osgeo import gdal
 from pytest_mock import MockerFixture
 
-from app.geospatial.wps_dataset import WPSDataset, multi_wps_dataset_context
+from app.geospatial.wps_dataset import WPSDataset
 from app.sfms.date_range_processor import BUIDateRangeProcessor
-from app.sfms.raster_addresser import RasterKeyAddresser
+from app.sfms.raster_addresser import FWIParameter, RasterKeyAddresser
 from app.tests.geospatial.test_wps_dataset import create_test_dataset
 from app.utils.s3_client import S3Client
 import uuid
 
 TEST_DATETIME = datetime(2024, 10, 10, 10, tzinfo=timezone.utc)
+EXPECTED_FIRST_DAY = TEST_DATETIME.replace(hour=20, minute=0, second=0, microsecond=0)
+EXPECTED_SECOND_DAY = TEST_DATETIME.replace(hour=20, minute=0, second=0, microsecond=0) + timedelta(days=1)
 
 
 def create_mock_gdal_dataset():
@@ -27,46 +29,48 @@ def create_mock_wps_dataset():
     return WPSDataset(ds=mock_ds, ds_path=None)
 
 
+def create_mock_wps_datasets(num: int) -> List[WPSDataset]:
+    return [create_mock_wps_dataset() for _ in range(num)]
+
+
 @pytest.mark.anyio
 async def test_bui_date_range_processor(mocker: MockerFixture):
     mock_key_addresser = RasterKeyAddresser()
+    get_weather_data_key_spy = mocker.spy(mock_key_addresser, "get_weather_data_keys")
+    gdal_prefix_keys_spy = mocker.spy(mock_key_addresser, "gdal_prefix_keys")
+    get_calculated_index_key_spy = mocker.spy(mock_key_addresser, "get_calculated_index_key")
     bui_date_range_processor = BUIDateRangeProcessor(TEST_DATETIME, 2, mock_key_addresser)
     # mock out storing of dataset
     mocker.patch.object(bui_date_range_processor, "_create_and_store_dataset", return_value="test_key.tif")
 
     # mock weather index, param datasets used for calculations
-    mock_dc_ds = create_mock_wps_dataset()
-    mock_dmc_ds = create_mock_wps_dataset()
-    mock_temp_ds = create_mock_wps_dataset()
-    mock_rh_ds = create_mock_wps_dataset()
-    mock_precip_ds = create_mock_wps_dataset()
+    input_datasets = create_mock_wps_datasets(5)
+    mock_temp_ds, mock_rh_ds, mock_precip_ds, mock_dc_ds, mock_dmc_ds = input_datasets
 
     @contextmanager
     def mock_input_dataset_context(_: List[str]):
-        mock_input_datasets = [mock_temp_ds, mock_rh_ds, mock_precip_ds, mock_dc_ds, mock_dmc_ds]
         try:
             # Enter each dataset's context and yield the list of instances
             with ExitStack() as stack:
-                yield [stack.enter_context(ds) for ds in mock_input_datasets]
+                yield [stack.enter_context(ds) for ds in input_datasets]
         finally:
             # Close all datasets to ensure cleanup
-            for ds in mock_input_datasets:
+            for ds in input_datasets:
                 ds.close()
 
     # mock new dmc and dc datasets
-    mock_new_dmc_ds = create_mock_wps_dataset()
-    mock_new_dc_ds = create_mock_wps_dataset()
+    new_datasets = create_mock_wps_datasets(2)
+    mock_new_dmc_ds, mock_new_dc_ds = new_datasets
 
     @contextmanager
     def mock_new_dmc_dc_datasets_context(_: List[str]):
-        mock_input_datasets = [mock_new_dmc_ds, mock_new_dc_ds]
         try:
             # Enter each dataset's context and yield the list of instances
             with ExitStack() as stack:
-                yield [stack.enter_context(ds) for ds in mock_input_datasets]
+                yield [stack.enter_context(ds) for ds in new_datasets]
         finally:
             # Close all datasets to ensure cleanup
-            for ds in mock_input_datasets:
+            for ds in new_datasets:
                 ds.close()
 
     # mock s3 client
@@ -78,4 +82,54 @@ def mock_new_dmc_dc_datasets_context(_: List[str]):
     mocker.patch("osgeo.gdal.Open", return_value=create_mock_gdal_dataset())
 
     await bui_date_range_processor.process_bui(mock_s3_client, mock_input_dataset_context, mock_new_dmc_dc_datasets_context)
-    mock_all_objects_exist.assert_called()
+
+    # Verify weather model keys and actual keys are checked for both days
+    assert mock_all_objects_exist.call_count == 4
+
+    # Verify the arguments for each call for get_weather_data_keys
+    assert get_weather_data_key_spy.call_args_list == [
+        mocker.call(TEST_DATETIME, EXPECTED_FIRST_DAY, 20),
+        mocker.call(TEST_DATETIME, EXPECTED_SECOND_DAY, 44),
+    ]
+
+    # Verify the arguments for each call for gdal_prefix_keys
+    assert gdal_prefix_keys_spy.call_args_list == [
+        # first day weather models
+        mocker.call(
+            "weather_models/rdps/2024-10-10/00/temp/CMC_reg_TMP_TGL_2_ps10km_2024101000_P020.grib2",
+            "weather_models/rdps/2024-10-10/00/rh/CMC_reg_RH_TGL_2_ps10km_2024101000_P020.grib2",
+            "weather_models/rdps/2024-10-10/12/precip/COMPUTED_reg_APCP_SFC_0_ps10km_20241010_20z.tif",
+        ),
+        # first day uploads
+        mocker.call("sfms/uploads/actual/2024-10-09/dc20241009.tif", "sfms/uploads/actual/2024-10-09/dmc20241009.tif"),
+        # second day weather models
+        mocker.call(
+            "weather_models/rdps/2024-10-10/00/temp/CMC_reg_TMP_TGL_2_ps10km_2024101000_P044.grib2",
+            "weather_models/rdps/2024-10-10/00/rh/CMC_reg_RH_TGL_2_ps10km_2024101000_P044.grib2",
+            "weather_models/rdps/2024-10-11/12/precip/COMPUTED_reg_APCP_SFC_0_ps10km_20241011_20z.tif",
+        ),
+        # second day uploads
+        mocker.call("sfms/calculated/forecast/2024-10-10/dc20241010.tif", "sfms/calculated/forecast/2024-10-10/dmc20241010.tif"),
+    ]
+
+    # Verify calculated keys are generated in order
+    assert get_calculated_index_key_spy.call_args_list == [
+        # first day
+        mocker.call(EXPECTED_FIRST_DAY, FWIParameter.DMC),
+        mocker.call(EXPECTED_FIRST_DAY, FWIParameter.DC),
+        mocker.call(EXPECTED_FIRST_DAY, FWIParameter.BUI),
+        # second day, previous days' dc and dmc are looked up
+        mocker.call(EXPECTED_FIRST_DAY, FWIParameter.DC),
+        mocker.call(EXPECTED_FIRST_DAY, FWIParameter.DMC),
+        mocker.call(EXPECTED_SECOND_DAY, FWIParameter.DMC),
+        mocker.call(EXPECTED_SECOND_DAY, FWIParameter.DC),
+        mocker.call(EXPECTED_SECOND_DAY, FWIParameter.BUI),
+    ]
+
+
+# 'weather_models/rdps/2024-10-10/00/temp/CMC_reg_TMP_TGL_2_ps10km_2024101000_P020.grib2'
+#'weather_models/rdps/2024-10-10/00/rh/CMC_reg_RH_TGL_2_ps10km_2024101000_P020.grib2'
+#'weather_models/rdps/2024-10-10/12/precip/COMPUTED_reg_APCP_SFC_0_ps10km_20241010_20z.tif'
+
+#'sfms/uploads/actual/2024-10-09/dc20241009.tif'
+# "sfms/uploads/actual/2024-10-09/dmc20241009.tif"