From 757e61c88b82e98f0cf93445cb48476edf1afefd Mon Sep 17 00:00:00 2001
From: Doug Latornell <dlatornell@eoas.ubc.ca>
Date: Mon, 4 Mar 2024 15:41:15 -0800
Subject: [PATCH] Automate creation of month-averaged datasets

Add launch of make_averaged_dataset worker for month-averaged datasets for
var_groups="biology", "chemistry", "physics" for "nowcast-green" after
completion of day-averaged dataset creation on the last day of the month. Also
added test cases for the functionality.
---
 nowcast/next_workers.py    | 35 ++++++++++++++++++++++++++++++++++-
 tests/test_next_workers.py | 30 ++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/nowcast/next_workers.py b/nowcast/next_workers.py
index 462b5d3a..b878e665 100644
--- a/nowcast/next_workers.py
+++ b/nowcast/next_workers.py
@@ -1548,7 +1548,40 @@ def after_make_averaged_dataset(msg, config, checklist):
     :returns: Worker(s) to launch next
     :rtype: list
     """
-    return []
+    next_workers = {
+        "crash": [],
+        "failure day biology": [],
+        "failure day chemistry": [],
+        "failure day physics": [],
+        "failure month biology": [],
+        "failure month chemistry": [],
+        "failure month physics": [],
+        "success day biology": [],
+        "success day chemistry": [],
+        "success day physics": [],
+        "success month biology": [],
+        "success month chemistry": [],
+        "success month physics": [],
+    }
+    if msg.type.startswith("success day"):
+        *_, reshapr_var_group = msg.type.split()
+        run_date = arrow.get(msg.payload[f"day {reshapr_var_group}"]["run date"])
+        if run_date.shift(days=+1).day == 1:
+            first_of_month = run_date.format("YYYY-MM-01")
+            next_workers[msg.type].append(
+                NextWorker(
+                    "nowcast.workers.make_averaged_dataset",
+                    args=[
+                        "skookum",
+                        "month",
+                        reshapr_var_group,
+                        "--run-date",
+                        first_of_month,
+                    ],
+                    host="localhost",
+                )
+            )
+    return next_workers[msg.type]
 
 
 def after_archive_tarball(msg, config, checklist):
diff --git a/tests/test_next_workers.py b/tests/test_next_workers.py
index 436a01c4..3c0ea34b 100644
--- a/tests/test_next_workers.py
+++ b/tests/test_next_workers.py
@@ -2221,6 +2221,36 @@ def test_no_next_worker_msg_types(self, msg_type, config, checklist):
         )
         assert workers == []
 
+    @pytest.mark.parametrize(
+        "msg_type",
+        [
+            "success day biology",
+            "success day chemistry",
+            "success day physics",
+        ],
+    )
+    def test_month_end_day_success_launch_month_average(
+        self, msg_type, config, checklist
+    ):
+        *_, reshapr_var_group = msg_type.split()
+        msg = Message(
+            "make_averaged_dataset",
+            msg_type,
+            payload={
+                f"day {reshapr_var_group}": {
+                    "run date": "2024-02-29",
+                    "file path": "SalishSea_1d_20240301_20240301_biol_T.nc",
+                }
+            },
+        )
+        workers = next_workers.after_make_averaged_dataset(msg, config, checklist)
+        expected = NextWorker(
+            "nowcast.workers.make_averaged_dataset",
+            args=["skookum", "month", reshapr_var_group, "--run-date", "2024-02-01"],
+            host="localhost",
+        )
+        assert expected in workers
+
 
 class TestAfterArchiveTarball:
     """Unit tests for the after_archive_tarball function."""