"""Unit tests for last_applied_treatment enrichment and product-JSON collapse.""" from __future__ import annotations import unittest from datetime import date from types import SimpleNamespace from unittest.mock import ANY, MagicMock, patch from pipeline.stages.advice_context import collapse_product_json, enrich_advice_json from pipeline.stages.treatments import load_last_treatment_before_recent_window from pipeline.stages.vector import RecommendedProduct _AS_OF = date(2026, 4, 16) _FIELD_ID = 7216 _SPRAY_DATE = date(2026, 4, 10) def _recommended(name: str) -> RecommendedProduct: return RecommendedProduct( product_name=name, registration_number=None, distance=None, query_index=0, product_id=name, ) class _StubBuilder: def build(self, conn, name): # noqa: ANN001 - matches ProductJsonBuilder.build return {name: {"uses": []}} def _payload(*, applied: list[str] | None = None) -> dict: return { "date_of_today": "16-04-2026", "meteorological_data": [ {"date": "11-04-2026", "applied_treatment": list(applied or [])}, ], "weather_forecasts": [], } def _enrich( advice_json: dict, *, recommended: list[RecommendedProduct] | None = None, last_advice_treatments: list[str] | None = None, last_advice_date: date | None = None, ) -> None: enrich_advice_json( advice_json, conn=MagicMock(), builder=_StubBuilder(), recommended=recommended or [], last_advice_summary="", last_advice_treatments=last_advice_treatments or [], as_of=_AS_OF, field_id=_FIELD_ID, last_advice_date=last_advice_date, ) class LastAppliedTreatmentTests(unittest.TestCase): def test_omits_when_allowed_products_empty_and_does_not_query(self) -> None: advice_json = _payload() with patch( "pipeline.stages.advice_context.load_last_treatment_before_recent_window" ) as loader: _enrich(advice_json, recommended=[]) loader.assert_not_called() self.assertNotIn("last_applied_treatment", advice_json) self.assertEqual(advice_json["allowed_products"], []) def test_omits_when_recent_applied_treatment_present(self) -> None: advice_json = _payload(applied=["LIETO"]) with patch( "pipeline.stages.advice_context.load_last_treatment_before_recent_window" ) as loader: _enrich(advice_json, recommended=[_recommended("LIETO")]) loader.assert_not_called() self.assertNotIn("last_applied_treatment", advice_json) def test_omits_when_loader_returns_none(self) -> None: advice_json = _payload() with patch( "pipeline.stages.advice_context.load_last_treatment_before_recent_window", return_value=None, ) as loader: _enrich(advice_json, recommended=[_recommended("LIETO")]) loader.assert_called_once_with(ANY, _FIELD_ID, _AS_OF) self.assertNotIn("last_applied_treatment", advice_json) def test_includes_at_end_when_all_gates_pass(self) -> None: advice_json = _payload() with patch( "pipeline.stages.advice_context.load_last_treatment_before_recent_window", return_value=(_SPRAY_DATE, ["LIETO"]), ): _enrich(advice_json, recommended=[_recommended("LIETO")]) self.assertEqual(list(advice_json)[-1], "last_applied_treatment") self.assertEqual( advice_json["last_applied_treatment"], {"date": "10-04-2026", "products": ["LIETO"]}, ) def test_name_only_when_product_in_allowed_products(self) -> None: advice_json = _payload() with patch( "pipeline.stages.advice_context.load_last_treatment_before_recent_window", return_value=(_SPRAY_DATE, ["LIETO"]), ): _enrich(advice_json, recommended=[_recommended("LIETO")]) self.assertEqual(advice_json["last_applied_treatment"]["products"], ["LIETO"]) def test_name_only_when_product_in_suggested_products_as_string(self) -> None: advice_json = _payload() with patch( "pipeline.stages.advice_context.load_last_treatment_before_recent_window", return_value=(_SPRAY_DATE, ["LIETO"]), ): _enrich( advice_json, recommended=[_recommended("LIETO")], last_advice_treatments=["LIETO"], ) self.assertEqual(advice_json["last_advice"]["suggested_products"], ["LIETO"]) self.assertEqual(advice_json["last_applied_treatment"]["products"], ["LIETO"]) def test_name_only_when_product_in_suggested_products_as_object(self) -> None: advice_json = _payload() with patch( "pipeline.stages.advice_context.load_last_treatment_before_recent_window", return_value=(_SPRAY_DATE, ["TWINGO"]), ): _enrich( advice_json, recommended=[_recommended("LIETO")], last_advice_treatments=["TWINGO"], ) self.assertEqual( advice_json["last_advice"]["suggested_products"], [{"TWINGO": {"uses": []}}], ) self.assertEqual(advice_json["last_applied_treatment"]["products"], ["TWINGO"]) def test_full_json_when_product_in_neither_list(self) -> None: advice_json = _payload() with patch( "pipeline.stages.advice_context.load_last_treatment_before_recent_window", return_value=(_SPRAY_DATE, ["UNKNOWN"]), ): _enrich(advice_json, recommended=[_recommended("LIETO")]) self.assertEqual( advice_json["last_applied_treatment"]["products"], [{"UNKNOWN": {"uses": []}}], ) def test_tank_mix_mixes_name_and_full_json(self) -> None: advice_json = _payload() with patch( "pipeline.stages.advice_context.load_last_treatment_before_recent_window", return_value=(_SPRAY_DATE, ["LIETO", "UNKNOWN"]), ): _enrich(advice_json, recommended=[_recommended("LIETO")]) self.assertEqual( advice_json["last_applied_treatment"]["products"], ["LIETO", {"UNKNOWN": {"uses": []}}], ) class CollapseProductJsonTests(unittest.TestCase): def test_collapses_last_applied_treatment_products(self) -> None: collapsed = collapse_product_json( { "allowed_products": [{"LIETO": {"uses": []}}], "last_advice": { "date": "15-04-2026", "advice_summary": "", "suggested_products": [{"TWINGO": {"uses": []}}], }, "last_applied_treatment": { "date": "10-04-2026", "products": [{"UNKNOWN": {"uses": []}}, "LIETO"], }, } ) self.assertEqual(collapsed["allowed_products"], ["LIETO"]) self.assertEqual(collapsed["last_advice"]["suggested_products"], ["TWINGO"]) self.assertEqual( collapsed["last_applied_treatment"], {"date": "10-04-2026", "products": ["UNKNOWN", "LIETO"]}, ) class LoadLastTreatmentBeforeRecentWindowTests(unittest.TestCase): def test_returns_latest_day_products(self) -> None: rows = [ SimpleNamespace(rilop_date=date(2026, 3, 22), prodgov_name="OLD"), SimpleNamespace(rilop_date=date(2026, 4, 10), prodgov_name="LIETO"), SimpleNamespace(rilop_date=date(2026, 4, 10), prodgov_name="TWINGO"), ] with patch("pipeline.stages.treatments.fetch_all", return_value=rows) as fetch: result = load_last_treatment_before_recent_window( MagicMock(), _FIELD_ID, _AS_OF ) params = fetch.call_args[0][2] self.assertEqual(params[1], date(2026, 3, 22)) # as_of-25 self.assertEqual(params[2], date(2026, 4, 10)) # as_of-6 self.assertEqual(result, (_SPRAY_DATE, ["LIETO", "TWINGO"])) def test_returns_none_when_no_rows(self) -> None: with patch("pipeline.stages.treatments.fetch_all", return_value=[]): result = load_last_treatment_before_recent_window( MagicMock(), _FIELD_ID, _AS_OF ) self.assertIsNone(result) if __name__ == "__main__": unittest.main()