From 47b4300391cef3c22649281d539d850b2c7566ab Mon Sep 17 00:00:00 2001
From: Arsham Mirehvandi <130353205+arshammi@users.noreply.github.com>
Date: Thu, 22 Oct 2026 09:35:16 +0200
Subject: [PATCH] Add query_forecast_end function and integrate into job
 processing

- Introduced the `query_forecast_end` function in `incubations.py` to determine the end date for query-synthesis weather based on incubation status and treatment intervals.
- Updated the `_run_job_inner` function in `job.py` to utilize `query_forecast_end`, enhancing the job's ability to handle future disease forecasting.
- Added unit tests in `test_incubations.py` to validate the behavior of `query_forecast_end` under various scenarios, ensuring accurate forecasting logic.
- Improved overall code clarity and maintainability with the new function integration.
---
 pipeline/job.py                |  3 ++-
 pipeline/stages/incubations.py | 21 ++++++++++++++++++
 tests/test_incubations.py      | 40 ++++++++++++++++++++++++++++++++++
 3 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/pipeline/job.py b/pipeline/job.py
index e71400c..38724cd 100644
--- a/pipeline/job.py
+++ b/pipeline/job.py
@@ -46,6 +46,7 @@ from pipeline.stages.incubations import (
     forecast_has_risk,
     incubation_start_covered,
     parse_register,
+    query_forecast_end,
     reconcile_register,
     register_payload,
 )
@@ -251,7 +252,7 @@ def _run_job_inner(resources: Resources, job: Job, as_of: date, dry_run: bool) -
         disease.disease_english,
         weather,
         phenology,
-        future_disease,
+        query_forecast_end(future_disease, incubation_register, as_of, interval_open),
     )
 
     payload: dict[str, Any] = {
diff --git a/pipeline/stages/incubations.py b/pipeline/stages/incubations.py
index 37cd312..0a6205a 100644
--- a/pipeline/stages/incubations.py
+++ b/pipeline/stages/incubations.py
@@ -214,6 +214,27 @@ def first_future_disease_date(
     return None
 
 
+def query_forecast_end(
+    future_disease: date | None,
+    register: dict[str, IncubationRecord],
+    as_of: date,
+    interval_open: bool,
+) -> date | None:
+    """End date for query-synthesis weather. None means today through as_of+5.
+
+    An uncovered incubation that has already started, with the treatment
+    interval closed, uses the full forward window. Otherwise the window ends
+    on the first future risk day.
+    """
+    underway = any(
+        not record.infection_covered and record.start_date <= as_of
+        for record in register.values()
+    )
+    if underway and not interval_open:
+        return None
+    return future_disease
+
+
 def covered_by_key(register: dict[str, IncubationRecord]) -> dict[str, bool]:
     return {key: record.infection_covered for key, record in register.items()}
 
diff --git a/tests/test_incubations.py b/tests/test_incubations.py
index 1228a77..b906a74 100644
--- a/tests/test_incubations.py
+++ b/tests/test_incubations.py
@@ -16,6 +16,7 @@ from pipeline.stages.incubations import (
     forecast_has_risk,
     incubation_start_covered,
     parse_register,
+    query_forecast_end,
     reconcile_register,
 )
 from pipeline.stages.products import ProductIndex, ProductRow
@@ -247,6 +248,45 @@ class RiskAndPrefilterTests(unittest.TestCase):
         register = {"incubazioneprimaria2": _record(_FUTURE, True)}
         self.assertEqual(first_future_disease_date(forecasts, register, _AS_OF), _FAR)
 
+    def test_underway_uncovered_incubation_uses_the_full_window(self) -> None:
+        tomorrow = _AS_OF + timedelta(days=1)
+        forecasts = {
+            _PAST: ["INCUBAZPRIMARIA2_40"],
+            tomorrow: ["INCUBAZPRIMARIA2_60"],
+            _FAR: ["FASE5"],
+        }
+        register = {"incubazioneprimaria2": _record(_PAST, False)}
+        future = first_future_disease_date(forecasts, register, _AS_OF)
+        self.assertEqual(future, tomorrow)
+        self.assertIsNone(query_forecast_end(future, register, _AS_OF, False))
+
+    def test_open_interval_keeps_the_first_future_day(self) -> None:
+        tomorrow = _AS_OF + timedelta(days=1)
+        forecasts = {tomorrow: ["INCUBAZPRIMARIA2_60"]}
+        register = {"incubazioneprimaria2": _record(_PAST, False)}
+        future = first_future_disease_date(forecasts, register, _AS_OF)
+        self.assertEqual(query_forecast_end(future, register, _AS_OF, True), tomorrow)
+
+    def test_future_only_incubation_keeps_its_start_day(self) -> None:
+        forecasts = {_FUTURE: ["INCUBAZPRIMARIA2_40"]}
+        register = {"incubazioneprimaria2": _record(_FUTURE, False)}
+        future = first_future_disease_date(forecasts, register, _AS_OF)
+        self.assertEqual(query_forecast_end(future, register, _AS_OF, False), _FUTURE)
+
+    def test_covered_incubation_keeps_a_later_fase5(self) -> None:
+        forecasts = {_PAST: ["INCUBAZPRIMARIA2_80"], _FAR: ["FASE5"]}
+        register = {"incubazioneprimaria2": _record(_PAST, True)}
+        future = first_future_disease_date(forecasts, register, _AS_OF)
+        self.assertEqual(future, _FAR)
+        self.assertEqual(query_forecast_end(future, register, _AS_OF, False), _FAR)
+
+    def test_incubation_starting_today_uses_the_full_window(self) -> None:
+        tomorrow = _AS_OF + timedelta(days=1)
+        forecasts = {_TODAY: ["INCUBAZPRIMARIA2_10"], tomorrow: ["INCUBAZPRIMARIA2_20"]}
+        register = {"incubazioneprimaria2": _record(_TODAY, False)}
+        future = first_future_disease_date(forecasts, register, _AS_OF)
+        self.assertIsNone(query_forecast_end(future, register, _AS_OF, False))
+
     def test_uncovered_past_incubation_selects_incubazprimaria(self) -> None:
         forecasts = {_PAST: ["INCUBAZPRIMARIA2_80.7", "FASE5"]}
         rule = select_prefilter_rule(forecasts, {}, _AS_OF, {"incubazioneprimaria2": False})