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.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}},
        )
        self.farm.plants.set([self.plant])

    @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_chat_client")
    def test_irrigation_recommendation_uses_farm_relations_when_request_omits_names(
        self,
        mock_get_chat_client,
        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_response = Mock()
        mock_response.choices = [Mock(message=Mock(content='{"plan": {"frequencyPerWeek": 2}}'))]
        mock_get_chat_client.return_value.chat.completions.create.return_value = mock_response

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

        self.assertEqual(result["plan"]["frequencyPerWeek"], 2)
        mock_build_rag_context.assert_called_once()
        mock_build_plant_text.assert_called_once_with("گوجه‌فرنگی", "میوه‌دهی")
        mock_build_irrigation_method_text.assert_called_once_with("آبیاری قطره‌ای")

    @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_chat_client")
    def test_fertilization_recommendation_uses_farm_plant_when_request_omits_name(
        self,
        mock_get_chat_client,
        _mock_build_rag_context,
        mock_build_plant_text,
    ):
        mock_response = Mock()
        mock_response.choices = [Mock(message=Mock(content='{"plan": {"npkRatio": "20-20-20"}}'))]
        mock_get_chat_client.return_value.chat.completions.create.return_value = mock_response

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

        self.assertEqual(result["plan"]["npkRatio"], "20-20-20")
        mock_build_plant_text.assert_called_once_with("گوجه‌فرنگی", "رویشی")