AI_Agro_Support/tests/test_assemble.py

"""Unit tests for advice-generation JSON assembly."""

from __future__ import annotations

import unittest
from datetime import date, timedelta
from typing import Any

from pipeline.stages.assemble import (
    assemble_json_for_advice_generation,
    build_json_for_query_synthesis,
)
from pipeline.window import build_window, format_date, today_and_future

_AS_OF = date(2026, 4, 16)
_PHASE = "BBCH 15"
_WIND_KEY = "average_wind_speed_m_s"


def _frc_weather(value: float | None) -> dict[date, dict[str, Any]]:
    return {day: {_WIND_KEY: value} for day in today_and_future(_AS_OF)}


class AdviceGenerationPhenologyTests(unittest.TestCase):
    def test_phenology_phase_only_on_past_and_today(self) -> None:
        window = build_window(_AS_OF)
        phenology = {day: _PHASE if day <= _AS_OF else None for day in window}
        payload = assemble_json_for_advice_generation(
            _AS_OF,
            weather={},
            phenology=phenology,
            forecasts={},
        )

        self.assertEqual(payload["date_of_today"], format_date(_AS_OF))
        self.assertEqual(
            list(payload),
            [
                "date_of_today",
                "predicted_incubations_management",
                "meteorological_data",
                "weather_forecasts",
            ],
        )
        self.assertEqual(payload["predicted_incubations_management"], {})
        self.assertEqual(len(payload["meteorological_data"]), 6)
        self.assertEqual(len(payload["weather_forecasts"]), 5)

        for entry in payload["meteorological_data"]:
            self.assertIn("phenology_phase", entry)
            self.assertEqual(entry["phenology_phase"], _PHASE)

        for entry in payload["weather_forecasts"]:
            self.assertNotIn("phenology_phase", entry)

        last_meteo = payload["meteorological_data"][-1]
        self.assertEqual(last_meteo["date"], format_date(_AS_OF))
        first_forecast = payload["weather_forecasts"][0]
        self.assertEqual(first_forecast["date"], format_date(_AS_OF + timedelta(days=1)))

    def test_register_precedes_weather_and_forecast_is_a_list(self) -> None:
        day = _AS_OF - timedelta(days=1)
        payload = assemble_json_for_advice_generation(
            _AS_OF,
            weather={},
            phenology={},
            forecasts={day: ["INCUBAZPRIMARIA2_80.7", "FASE5"]},
            incubation_register={
                "incubazioneprimaria2": {
                    "start_date": "15-04-2026",
                    "infection_covered": False,
                }
            },
        )
        self.assertEqual(
            list(payload)[:3],
            ["date_of_today", "predicted_incubations_management", "meteorological_data"],
        )
        self.assertEqual(
            payload["predicted_incubations_management"]["incubazioneprimaria2"]["infection_covered"],
            False,
        )
        match = next(entry for entry in payload["meteorological_data"] if entry["date"] == format_date(day))
        self.assertEqual(match["disease_forecast"], ["INCUBAZPRIMARIA2_80.7", "FASE5"])


class ForecastWindSpeedTests(unittest.TestCase):
    def test_wind_on_today_and_future_not_past(self) -> None:
        payload = assemble_json_for_advice_generation(
            _AS_OF,
            weather=_frc_weather(1.234),
            phenology={},
            forecasts={},
        )

        for entry in payload["meteorological_data"][:-1]:
            self.assertNotIn(_WIND_KEY, entry)

        today = payload["meteorological_data"][-1]
        self.assertEqual(today["date"], format_date(_AS_OF))
        self.assertEqual(today[_WIND_KEY], 1.23)

        for entry in payload["weather_forecasts"]:
            self.assertEqual(entry[_WIND_KEY], 1.23)

    def test_missing_frc_wind_is_null(self) -> None:
        payload = assemble_json_for_advice_generation(
            _AS_OF,
            weather=_frc_weather(None),
            phenology={},
            forecasts={},
        )

        for entry in payload["meteorological_data"][:-1]:
            self.assertNotIn(_WIND_KEY, entry)

        self.assertIsNone(payload["meteorological_data"][-1][_WIND_KEY])
        for entry in payload["weather_forecasts"]:
            self.assertIn(_WIND_KEY, entry)
            self.assertIsNone(entry[_WIND_KEY])

    def test_query_synthesis_includes_wind_on_every_date(self) -> None:
        payload = build_json_for_query_synthesis(
            _AS_OF,
            "grapevine",
            "downy mildew",
            _frc_weather(2.567),
            {_AS_OF: _PHASE},
            None,
        )

        self.assertEqual(len(payload["weather_forecasts"]), 6)
        self.assertEqual(payload["weather_forecasts"][0]["date"], format_date(_AS_OF))
        for entry in payload["weather_forecasts"]:
            self.assertEqual(entry[_WIND_KEY], 2.57)

    def test_query_synthesis_missing_wind_is_null(self) -> None:
        weather = {
            day: {_WIND_KEY: None}
            for day in today_and_future(_AS_OF)
            if day <= _AS_OF + timedelta(days=2)
        }
        payload = build_json_for_query_synthesis(
            _AS_OF,
            "grapevine",
            "downy mildew",
            weather,
            {},
            _AS_OF + timedelta(days=2),
        )

        self.assertEqual(len(payload["weather_forecasts"]), 3)
        for entry in payload["weather_forecasts"]:
            self.assertIn(_WIND_KEY, entry)
            self.assertIsNone(entry[_WIND_KEY])


if __name__ == "__main__":
    unittest.main()