"""Unit tests for last_applied_treatment enrichment and product-JSON collapse.""" from __future__ import annotations import unittest from datetime import date from unittest.mock import MagicMock, patch from pipeline.stages.advice_context import collapse_product_json, enrich_advice_json from pipeline.stages.treatments import TreatmentHistory from pipeline.stages.vector import RecommendedProduct _AS_OF = date(2026, 4, 16) _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 _history(*product_names: str, spray_date: date = _SPRAY_DATE) -> TreatmentHistory: by_day = {spray_date: list(product_names)} if product_names else {} return TreatmentHistory(as_of=_AS_OF, by_day=by_day) def _enrich( advice_json: dict, *, recommended: list[RecommendedProduct] | None = None, last_advice_treatments: list[str] | None = None, last_advice_date: date | None = None, treatment_history: TreatmentHistory | 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 [], last_advice_date=last_advice_date, treatment_history=treatment_history if treatment_history is not None else _history(), ) class LastAppliedTreatmentTests(unittest.TestCase): def test_omits_when_allowed_products_empty_and_does_not_query(self) -> None: advice_json = _payload() with patch("pipeline.stages.treatments.fetch_all") as fetch: _enrich( advice_json, recommended=[], treatment_history=_history("LIETO"), ) fetch.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.treatments.fetch_all") as fetch: _enrich( advice_json, recommended=[_recommended("LIETO")], treatment_history=_history("LIETO"), ) fetch.assert_not_called() self.assertNotIn("last_applied_treatment", advice_json) def test_omits_when_history_has_no_older_spray(self) -> None: advice_json = _payload() with patch("pipeline.stages.treatments.fetch_all") as fetch: _enrich( advice_json, recommended=[_recommended("LIETO")], treatment_history=TreatmentHistory(as_of=_AS_OF, by_day={}), ) fetch.assert_not_called() self.assertNotIn("last_applied_treatment", advice_json) def test_includes_at_end_when_all_gates_pass(self) -> None: advice_json = _payload() _enrich( advice_json, recommended=[_recommended("LIETO")], treatment_history=_history("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() _enrich( advice_json, recommended=[_recommended("LIETO")], treatment_history=_history("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() _enrich( advice_json, recommended=[_recommended("LIETO")], last_advice_treatments=["LIETO"], treatment_history=_history("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() _enrich( advice_json, recommended=[_recommended("LIETO")], last_advice_treatments=["TWINGO"], treatment_history=_history("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() _enrich( advice_json, recommended=[_recommended("LIETO")], treatment_history=_history("UNKNOWN"), ) self.assertEqual( advice_json["last_applied_treatment"]["products"], [{"UNKNOWN": {"uses": []}}], ) def test_tank_mix_mixes_name_and_full_json(self) -> None: advice_json = _payload() _enrich( advice_json, recommended=[_recommended("LIETO")], treatment_history=_history("LIETO", "UNKNOWN"), ) self.assertEqual( advice_json["last_applied_treatment"]["products"], ["LIETO", {"UNKNOWN": {"uses": []}}], ) def test_does_not_query_operations_when_including_last_applied_treatment(self) -> None: advice_json = _payload() with patch("pipeline.stages.treatments.fetch_all") as fetch: _enrich( advice_json, recommended=[_recommended("LIETO")], treatment_history=_history("LIETO"), ) fetch.assert_not_called() self.assertIn("last_applied_treatment", advice_json) 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"]}, ) if __name__ == "__main__": unittest.main()