From fd672105f280548999abf4fae7a08c2947f78bc6 Mon Sep 17 00:00:00 2001
From: Doug Latornell <dlatornell@eoas.ubc.ca>
Date: Thu, 26 Sep 2024 11:56:49 -0700
Subject: [PATCH] Add month-avg grazing and growth to automation

Expanded the next_workers module to calculate month-average grazing and growth
datasets at month-end. Added corresponding test functions to ensure the new
make_averaged_dataset worker instances are launched correctly.
---
 nowcast/next_workers.py    | 31 ++++++++++++++++++++++++++++
 tests/test_next_workers.py | 42 +++++++++++++++++++++++++++++++++++++-
 2 files changed, 72 insertions(+), 1 deletion(-)

diff --git a/nowcast/next_workers.py b/nowcast/next_workers.py
index 6fb0d59f..7e68c66d 100644
--- a/nowcast/next_workers.py
+++ b/nowcast/next_workers.py
@@ -1557,6 +1557,8 @@ def after_make_averaged_dataset(msg, config, checklist):
         "crash": [],
         "failure day biology": [],
         "failure day chemistry": [],
+        "failure day grazing": [],
+        "failure day growth": [],
         "failure day physics": [],
         "failure month biology": [],
         "failure month chemistry": [],
@@ -1566,6 +1568,8 @@ def after_make_averaged_dataset(msg, config, checklist):
         "success day physics": [],
         "success month biology": [],
         "success month chemistry": [],
+        "success month grazing": [],
+        "success month growth": [],
         "success month physics": [],
     }
     if msg.type.startswith("success day"):
@@ -1580,6 +1584,33 @@ def after_make_averaged_dataset(msg, config, checklist):
                     host="localhost",
                 )
             )
+    if msg.type.startswith("success month"):
+        *_, reshapr_var_group = msg.type.split()
+        match reshapr_var_group:
+            case "physics":
+                run_date = arrow.get(msg.payload["month physics"]["run date"]).format(
+                    "YYYY-MM-DD"
+                )
+                next_workers[msg.type].append(
+                    NextWorker(
+                        "nowcast.workers.make_averaged_dataset",
+                        args=["month", "grazing", "--run-date", run_date],
+                        host="localhost",
+                    )
+                )
+            case "grazing":
+                run_date = arrow.get(msg.payload["month grazing"]["run date"]).format(
+                    "YYYY-MM-DD"
+                )
+                next_workers[msg.type].append(
+                    NextWorker(
+                        "nowcast.workers.make_averaged_dataset",
+                        args=["month", "growth", "--run-date", run_date],
+                        host="localhost",
+                    )
+                )
+            case _:
+                pass
     return next_workers[msg.type]
 
 
diff --git a/tests/test_next_workers.py b/tests/test_next_workers.py
index d5589997..e62ba773 100644
--- a/tests/test_next_workers.py
+++ b/tests/test_next_workers.py
@@ -2216,10 +2216,12 @@ class TestAfterMakeAveragedDataset:
             "crash",
             "failure day biology",
             "failure day chemistry",
+            "failure day grazing",
+            "failure day growth",
             "failure day physics",
             "success month biology",
             "success month chemistry",
-            "success month physics",
+            "success month growth",
             "failure month biology",
             "failure month chemistry",
             "failure month physics",
@@ -2261,6 +2263,44 @@ def test_month_end_day_success_launch_month_average(
         )
         assert expected in workers
 
+    def test_month_physics_success_launch_month_grazing(self, config, checklist):
+        msg = Message(
+            "make_averaged_dataset",
+            "success month physics",
+            payload={
+                "month physics": {
+                    "run date": "2024-09-01",
+                    "file path": "SalishSea_1m_20240901_20240930_grid_T.nc",
+                }
+            },
+        )
+        workers = next_workers.after_make_averaged_dataset(msg, config, checklist)
+        expected = NextWorker(
+            "nowcast.workers.make_averaged_dataset",
+            args=["month", "grazing", "--run-date", "2024-09-01"],
+            host="localhost",
+        )
+        assert expected in workers
+
+    def test_month_grazing_success_launch_month_growth(self, config, checklist):
+        msg = Message(
+            "make_averaged_dataset",
+            "success month grazing",
+            payload={
+                "month grazing": {
+                    "run date": "2024-09-01",
+                    "file path": "SalishSea_1m_20240901_20240930_graz_T.nc",
+                }
+            },
+        )
+        workers = next_workers.after_make_averaged_dataset(msg, config, checklist)
+        expected = NextWorker(
+            "nowcast.workers.make_averaged_dataset",
+            args=["month", "growth", "--run-date", "2024-09-01"],
+            host="localhost",
+        )
+        assert expected in workers
+
 
 class TestAfterArchiveTarball:
     """Unit tests for the after_archive_tarball function."""