AI_Agro_Support/tests/test_advice_context.py

"""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_includes_older_spray_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.assertEqual(advice_json["allowed_products"], [])
        self.assertEqual(
            advice_json["last_applied_treatment"],
            {"date": "10-04-2026", "products": [{"LIETO": {"uses": []}}]},
        )

    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()