Ai/rag/tests/test_recommendation_services.py

import uuid
from datetime import date
from unittest.mock import Mock, patch

from django.test import TestCase

from farm_data.models import SensorData
from irrigation.models import IrrigationMethod
from location_data.models import SoilLocation
from plant.models import Plant
from rag.services.fertilization import get_fertilization_recommendation
from rag.services.irrigation import get_irrigation_recommendation
from weather.models import WeatherForecast


class RecommendationServiceDefaultsTests(TestCase):
    def setUp(self):
        self.location = SoilLocation.objects.create(
            latitude="35.700000",
            longitude="51.400000",
            farm_boundary={"type": "Polygon", "coordinates": []},
        )
        WeatherForecast.objects.create(
            location=self.location,
            forecast_date=date(2026, 4, 10),
            temperature_min=12.0,
            temperature_max=23.0,
            temperature_mean=18.0,
        )
        self.plant = Plant.objects.create(name="گوجه‌فرنگی")
        self.onion = Plant.objects.create(name="پیاز")
        self.irrigation_method = IrrigationMethod.objects.create(name="آبیاری قطره‌ای")
        self.farm_uuid = uuid.uuid4()
        self.farm = SensorData.objects.create(
            farm_uuid=self.farm_uuid,
            center_location=self.location,
            irrigation_method=self.irrigation_method,
            sensor_payload={
                "sensor-7-1": {
                    "soil_moisture": 30.0,
                    "nitrogen": 18.0,
                    "phosphorus": 12.0,
                    "potassium": 14.0,
                    "soil_ph": 6.9,
                }
            },
        )
        self.farm.plants.set([self.plant])

    def build_irrigation_optimizer_result(self):
        return {
            "engine": "crop_simulation_heuristic",
            "context_text": "optimizer irrigation context",
            "recommended_strategy": {
                "code": "balanced",
                "label": "آبیاری متعادل",
                "score": 88.0,
                "expected_yield_index": 91.0,
                "total_irrigation_mm": 24.0,
                "amount_per_event_mm": 8.0,
                "events": 3,
                "frequency_per_week": 3,
                "event_dates": ["2026-04-10"],
                "timing": "اوایل صبح",
                "moisture_target_percent": 70,
                "validity_period": "معتبر برای 3 روز آینده",
                "reasoning": ["شبیه ساز این سناریو را برتر ارزیابی کرد."],
            },
            "alternatives": [
                {
                    "code": "protective",
                    "label": "آبیاری حمایتی",
                "score": 80.0,
                "expected_yield_index": 85.0,
                "total_irrigation_mm": 28.0,
                }
            ],
        }

    def build_irrigation_llm_result(self):
        return (
            '{"plan": {"frequencyPerWeek": 3, "durationMinutes": 42, "bestTimeOfDay": "اوایل صبح", '
            '"moistureLevel": 68, "warning": "بررسی شود"}, '
            '"timeline": [{"step_number": 1, "title": "بازبینی", "description": "لاین ها بررسی شوند"}], '
            '"sections": [{"type": "warning", "title": "هشدار", "icon": "alert-triangle", "content": "بررسی شود"}, '
            '{"type": "tip", "title": "نکته", "icon": "bulb", "content": "مورد سفارشی"}]}'
        )

    def build_fertilization_optimizer_result(self):
        return {
            "engine": "crop_simulation_heuristic",
            "context_text": "optimizer fertilization context",
            "recommended_strategy": {
                "code": "balanced",
                "label": "تغذیه متعادل",
                "score": 84.0,
                "expected_yield_index": 88.0,
                "fertilizer_type": "20-20-20",
                "amount_kg_per_ha": 65.0,
                "application_method": "کودآبیاری",
                "timing": "صبح زود",
                "validity_period": "معتبر برای 7 روز آینده",
                "reasoning": ["کسری عناصر با این سناریو بهتر پوشش داده می شود."],
            },
            "alternatives": [
                {
                    "code": "maintenance",
                    "label": "تغذیه نگهدارنده",
                    "score": 72.0,
                    "expected_yield_index": 78.0,
                    "fertilizer_type": "20-20-20",
                    "amount_kg_per_ha": 45.0,
                    "application_method": "کودآبیاری",
                    "timing": "صبح زود",
                    "reasoning": ["برای نگهداری تعادل تغذیه ای گزینه سبک تری است."],
                }
            ],
        }

    @patch("rag.services.irrigation.calculate_forecast_water_needs", return_value=[])
    @patch("rag.services.irrigation.resolve_kc", return_value=0.9)
    @patch("rag.services.irrigation.resolve_crop_profile", return_value={})
    @patch("rag.services.irrigation.build_irrigation_method_text", return_value="method text")
    @patch("rag.services.irrigation.build_plant_text", return_value="plant text")
    @patch("rag.services.irrigation.build_rag_context", return_value="")
    @patch("rag.services.irrigation._get_optimizer")
    @patch("rag.services.irrigation.get_chat_client")
    def test_irrigation_recommendation_uses_farm_relations_when_request_omits_names(
        self,
        mock_get_chat_client,
        mock_get_optimizer,
        mock_build_rag_context,
        mock_build_plant_text,
        mock_build_irrigation_method_text,
        _mock_resolve_crop_profile,
        _mock_resolve_kc,
        _mock_calculate_forecast_water_needs,
    ):
        mock_get_optimizer.return_value.optimize_irrigation.return_value = (
            self.build_irrigation_optimizer_result()
        )
        mock_response = Mock()
        mock_response.choices = [Mock(message=Mock(content=self.build_irrigation_llm_result()))]
        mock_get_chat_client.return_value.chat.completions.create.return_value = mock_response

        result = get_irrigation_recommendation(
            farm_uuid=str(self.farm_uuid),
            growth_stage="میوه‌دهی",
        )

        self.assertEqual(result["plan"]["frequencyPerWeek"], 3)
        self.assertEqual(result["plan"]["bestTimeOfDay"], "اوایل صبح")
        mock_build_rag_context.assert_called_once()
        mock_build_plant_text.assert_called_once_with("گوجه‌فرنگی", "میوه‌دهی")
        mock_build_irrigation_method_text.assert_called_once_with("آبیاری قطره‌ای")
        self.assertEqual(
            result["selected_irrigation_method"]["name"],
            "آبیاری قطره‌ای",
        )
        self.assertEqual(result["simulation_optimizer"]["engine"], "crop_simulation_heuristic")
        self.assertEqual(result["timeline"][0]["title"], "بازبینی")
        self.assertEqual(result["sections"][1]["type"], "tip")
        self.assertEqual(result["water_balance"]["active_kc"], 0.9)

    @patch("rag.services.irrigation.calculate_forecast_water_needs", return_value=[])
    @patch("rag.services.irrigation.resolve_kc", return_value=0.9)
    @patch("rag.services.irrigation.resolve_crop_profile", return_value={})
    @patch("rag.services.irrigation.build_irrigation_method_text", return_value="method text")
    @patch("rag.services.irrigation.build_plant_text", return_value="plant text")
    @patch("rag.services.irrigation.build_rag_context", return_value="")
    @patch("rag.services.irrigation._get_optimizer")
    @patch("rag.services.irrigation.get_chat_client")
    def test_irrigation_recommendation_persists_selected_method_on_farm(
        self,
        mock_get_chat_client,
        mock_get_optimizer,
        _mock_build_rag_context,
        _mock_build_plant_text,
        mock_build_irrigation_method_text,
        _mock_resolve_crop_profile,
        _mock_resolve_kc,
        _mock_calculate_forecast_water_needs,
    ):
        sprinkler = IrrigationMethod.objects.create(name="بارانی")
        self.farm.irrigation_method = None
        self.farm.save(update_fields=["irrigation_method", "updated_at"])

        mock_get_optimizer.return_value.optimize_irrigation.return_value = (
            self.build_irrigation_optimizer_result()
        )
        mock_response = Mock()
        mock_response.choices = [Mock(message=Mock(content=self.build_irrigation_llm_result()))]
        mock_get_chat_client.return_value.chat.completions.create.return_value = mock_response

        result = get_irrigation_recommendation(
            farm_uuid=str(self.farm_uuid),
            growth_stage="میوه‌دهی",
            irrigation_method_name="بارانی",
        )

        self.farm.refresh_from_db()
        self.assertEqual(self.farm.irrigation_method_id, sprinkler.id)
        self.assertEqual(result["selected_irrigation_method"]["id"], sprinkler.id)
        mock_build_irrigation_method_text.assert_called_once_with("بارانی")
        self.assertEqual(result["plan"]["warning"], "بررسی شود")

    @patch("rag.services.irrigation.calculate_forecast_water_needs", return_value=[])
    @patch("rag.services.irrigation.resolve_kc", return_value=0.9)
    @patch("rag.services.irrigation.resolve_crop_profile", return_value={})
    @patch("rag.services.irrigation.build_irrigation_method_text", return_value="method text")
    @patch("rag.services.irrigation.build_plant_text", return_value="plant text")
    @patch("rag.services.irrigation.build_rag_context", return_value="")
    @patch("rag.services.irrigation._get_optimizer")
    @patch("rag.services.irrigation.get_chat_client")
    def test_irrigation_recommendation_falls_back_to_optimizer_when_llm_returns_invalid_payload(
        self,
        mock_get_chat_client,
        mock_get_optimizer,
        _mock_build_rag_context,
        _mock_build_plant_text,
        _mock_build_irrigation_method_text,
        _mock_resolve_crop_profile,
        _mock_resolve_kc,
        _mock_calculate_forecast_water_needs,
    ):
        mock_get_optimizer.return_value.optimize_irrigation.return_value = (
            self.build_irrigation_optimizer_result()
        )
        mock_response = Mock()
        mock_response.choices = [Mock(message=Mock(content="not-json"))]
        mock_get_chat_client.return_value.chat.completions.create.return_value = mock_response

        result = get_irrigation_recommendation(
            farm_uuid=str(self.farm_uuid),
            growth_stage="میوه‌دهی",
        )

        self.assertEqual(result["plan"]["frequencyPerWeek"], 3)
        self.assertEqual(result["timeline"][0]["step_number"], 1)
        self.assertEqual(result["sections"][0]["type"], "warning")
        self.assertEqual(result["simulation_optimizer"]["engine"], "crop_simulation_heuristic")

    @patch("rag.services.fertilization.build_plant_text", return_value="plant text")
    @patch("rag.services.fertilization.build_rag_context", return_value="")
    @patch("rag.services.fertilization._get_optimizer")
    @patch("rag.services.fertilization.get_chat_client")
    def test_fertilization_recommendation_uses_farm_plant_when_request_omits_name(
        self,
        mock_get_chat_client,
        mock_get_optimizer,
        _mock_build_rag_context,
        mock_build_plant_text,
    ):
        mock_get_optimizer.return_value.optimize_fertilization.return_value = (
            self.build_fertilization_optimizer_result()
        )
        mock_response = Mock()
        mock_response.choices = [Mock(message=Mock(content='{"status": "success", "data": {"primary_recommendation": {"display_title": "کود کامل 20-20-20", "reasoning": "توضیح", "summary": "مصرف انجام شود"}, "nutrient_analysis": {"macro": [{"key": "n", "name": "نیتروژن (N)", "value": 20, "unit": "percent", "description": "..."}, {"key": "p", "name": "فسفر (P)", "value": 20, "unit": "percent", "description": "..."}, {"key": "k", "name": "پتاسیم (K)", "value": 20, "unit": "percent", "description": "..."}], "micro": []}, "application_guide": {"safety_warning": "از اختلاط نامناسب خودداری شود.", "steps": [{"step_number": 1, "title": "آماده سازی", "description": "مورد 1"}]}, "alternative_recommendations": []}}'))]
        mock_get_chat_client.return_value.chat.completions.create.return_value = mock_response

        result = get_fertilization_recommendation(
            farm_uuid=str(self.farm_uuid),
            growth_stage="رویشی",
        )

        self.assertEqual(result["data"]["primary_recommendation"]["npk_ratio"]["label"], "20-20-20")
        self.assertEqual(result["data"]["primary_recommendation"]["dosage"]["base_amount_per_hectare"], 65.0)
        mock_build_plant_text.assert_called_once_with("گوجه‌فرنگی", "رویشی")
        self.assertEqual(result["simulation_optimizer"]["engine"], "crop_simulation_heuristic")
        self.assertEqual(result["data"]["application_guide"]["safety_warning"], "از اختلاط نامناسب خودداری شود.")

    @patch("rag.services.fertilization.build_plant_text", return_value="plant text")
    @patch("rag.services.fertilization.build_rag_context", return_value="")
    @patch("rag.services.fertilization._get_optimizer")
    @patch("rag.services.fertilization.get_chat_client")
    def test_fertilization_recommendation_resolves_requested_plant_from_catalog(
        self,
        mock_get_chat_client,
        mock_get_optimizer,
        _mock_build_rag_context,
        mock_build_plant_text,
    ):
        mock_get_optimizer.return_value.optimize_fertilization.return_value = (
            self.build_fertilization_optimizer_result()
        )
        mock_response = Mock()
        mock_response.choices = [Mock(message=Mock(content="not-json"))]
        mock_get_chat_client.return_value.chat.completions.create.return_value = mock_response

        result = get_fertilization_recommendation(
            farm_uuid=str(self.farm_uuid),
            plant_name="پیاز",
            growth_stage="گلدهی",
        )

        optimizer_call = mock_get_optimizer.return_value.optimize_fertilization.call_args.kwargs
        self.assertEqual(getattr(optimizer_call["plant"], "name", None), "پیاز")
        self.assertEqual(optimizer_call["growth_stage"], "flowering")
        mock_build_plant_text.assert_called_once_with("پیاز", "flowering")
        self.assertEqual(result["data"]["primary_recommendation"]["npk_ratio"]["label"], "20-20-20")

    @patch("rag.services.fertilization.build_plant_text", return_value="plant text")
    @patch("rag.services.fertilization.build_rag_context", return_value="")
    @patch("rag.services.fertilization._get_optimizer")
    @patch("rag.services.fertilization.get_chat_client")
    def test_fertilization_recommendation_falls_back_to_optimizer_when_llm_returns_invalid_payload(
        self,
        mock_get_chat_client,
        mock_get_optimizer,
        _mock_build_rag_context,
        _mock_build_plant_text,
    ):
        mock_get_optimizer.return_value.optimize_fertilization.return_value = (
            self.build_fertilization_optimizer_result()
        )
        mock_response = Mock()
        mock_response.choices = [Mock(message=Mock(content="not-json"))]
        mock_get_chat_client.return_value.chat.completions.create.return_value = mock_response

        result = get_fertilization_recommendation(
            farm_uuid=str(self.farm_uuid),
            growth_stage="رویشی",
        )

        self.assertEqual(result["data"]["primary_recommendation"]["npk_ratio"]["label"], "20-20-20")
        self.assertEqual(result["data"]["primary_recommendation"]["dosage"]["base_amount_per_hectare"], 65.0)