Ai/crop_simulation/tests.py

import importlib.util
import os
import sqlite3
import tempfile
from collections import namedtuple
from datetime import date, timedelta
from unittest.mock import patch
from unittest import skipUnless

from django.test import TestCase

from .models import SimulationRun, SimulationScenario
from .services import CropSimulationService, CropSimulationError, PcseSimulationManager


def build_weather(days: int = 5) -> list[dict]:
    start = date(2026, 4, 1)
    return [
        {
            "DAY": start + timedelta(days=index),
            "LAT": 35.7,
            "LON": 51.4,
            "ELEV": 1200,
            "IRRAD": 16_000_000 + (index * 100_000),
            "TMIN": 11 + index,
            "TMAX": 22 + index,
            "VAP": 12,
            "WIND": 2.4,
            "RAIN": 0.8,
            "E0": 0.35,
            "ES0": 0.3,
            "ET0": 0.32,
        }
        for index in range(days)
    ]


def build_agromanagement(n_amount: float = 30.0) -> list[dict]:
    return [
        {
            date(2026, 4, 1): {
                "CropCalendar": {
                    "crop_name": "wheat",
                    "variety_name": "winter-wheat",
                    "crop_start_date": date(2026, 4, 5),
                    "crop_start_type": "sowing",
                    "crop_end_date": date(2026, 9, 1),
                    "crop_end_type": "harvest",
                    "max_duration": 180,
                },
                "TimedEvents": [
                    {
                        "event_signal": "apply_n",
                        "name": "N strategy",
                        "events_table": [
                            {
                                date(2026, 4, 20): {
                                    "N_amount": n_amount,
                                    "N_recovery": 0.7,
                                }
                            }
                        ],
                    }
                ],
                "StateEvents": [],
            }
        },
        {},
    ]


class CropSimulationServiceTests(TestCase):
    def setUp(self):
        self.service = CropSimulationService()
        self.weather = build_weather()
        self.soil = {"SMFCF": 0.34, "SMW": 0.12, "RDMSOL": 120.0}
        self.site = {"WAV": 40.0}
        self.crop = {"crop_name": "wheat", "TSUM1": 800, "YIELD_SCALE": 1.0}

    def test_failure_marks_scenario_and_run_failed(self):
        with patch.object(
            self.service.manager,
            "run_simulation",
            side_effect=CropSimulationError("pcse failed"),
        ):
            with self.assertRaises(CropSimulationError):
                self.service.run_single_simulation(
                    weather=self.weather,
                    soil=self.soil,
                    crop_parameters=self.crop,
                    agromanagement=build_agromanagement(),
                    site_parameters=self.site,
                    name="broken run",
                )

        scenario = SimulationScenario.objects.get()
        run = SimulationRun.objects.get()

        self.assertEqual(scenario.status, SimulationScenario.Status.FAILURE)
        self.assertEqual(run.status, SimulationScenario.Status.FAILURE)
        self.assertEqual(scenario.error_message, "pcse failed")

    def test_requires_at_least_two_fertilization_strategies(self):
        with self.assertRaises(CropSimulationError):
            self.service.compare_fertilization_strategies(
                weather=self.weather,
                soil=self.soil,
                crop_parameters=self.crop,
                strategies=[{"label": "only", "agromanagement": build_agromanagement()}],
                site_parameters=self.site,
            )

    def test_raises_clear_error_when_pcse_is_unavailable(self):
        with patch("crop_simulation.services._load_pcse_bindings", return_value=None):
            with self.assertRaisesMessage(
                CropSimulationError,
                "PCSE is not installed or required PCSE classes could not be loaded.",
            ):
                self.service.run_single_simulation(
                    weather=self.weather,
                    soil=self.soil,
                    crop_parameters=self.crop,
                    agromanagement=build_agromanagement(),
                    site_parameters=self.site,
                    name="missing pcse",
                )


@skipUnless(
    importlib.util.find_spec("pcse") is not None,
    "pcse must be installed to run real WOFOST integration tests.",
)
class CropSimulationPcseIntegrationTests(TestCase):
    def setUp(self):
        os.environ["HOME"] = tempfile.mkdtemp(prefix="pcse-home-", dir="/tmp")
        from pcse import settings as pcse_settings
        from pcse.tests.db_input import (
            AgroManagementDataProvider,
            GridWeatherDataProvider,
            fetch_cropdata,
            fetch_sitedata,
            fetch_soildata,
        )

        def namedtuple_factory(cursor, row):
            fields = [column[0] for column in cursor.description]
            cls = namedtuple("Row", fields)
            return cls._make(row)

        self.grid = int(os.environ.get("PCSE_TEST_GRID", "31031"))
        self.crop_no = int(os.environ.get("PCSE_TEST_CROP_NO", "1"))
        self.year = int(os.environ.get("PCSE_TEST_YEAR", "2000"))

        self.connection = sqlite3.connect(
            os.path.join(pcse_settings.PCSE_USER_HOME, "pcse.db")
        )
        self.connection.row_factory = namedtuple_factory

        self.weather = GridWeatherDataProvider(
            self.connection,
            grid_no=self.grid,
        ).export()
        self.soil = fetch_soildata(self.connection, self.grid)
        self.site = fetch_sitedata(self.connection, self.grid, self.year)
        self.crop = fetch_cropdata(
            self.connection,
            self.grid,
            self.year,
            self.crop_no,
        )
        self.agromanagement = AgroManagementDataProvider(
            self.connection,
            self.grid,
            self.crop_no,
            self.year,
        )
        self.service = CropSimulationService(
            manager=PcseSimulationManager(model_name="Wofost72_WLP_CWB")
        )

    def tearDown(self):
        self.connection.close()

    def test_real_wofost_execute_full_service_path(self):
        result = self.service.run_single_simulation(
            weather=self.weather,
            soil=self.soil,
            crop_parameters=self.crop,
            agromanagement=self.agromanagement,
            site_parameters=self.site,
            name="pcse path",
        )

        scenario = SimulationScenario.objects.get()

        self.assertEqual(scenario.status, SimulationScenario.Status.SUCCESS)
        self.assertEqual(result["result"]["engine"], "pcse")
        self.assertIsNotNone(result["result"]["metrics"]["yield_estimate"])
        self.assertIsNotNone(result["result"]["metrics"]["biomass"])