Ai/rag/tests/test_recommendation_api.py

from unittest.mock import patch

from django.test import TestCase, override_settings
from rest_framework.test import APIClient


@override_settings(ROOT_URLCONF="config.test_urls")
class RagRecommendationApiTests(TestCase):
    def setUp(self):
        self.client = APIClient()

    @patch("rag.services.irrigation.get_irrigation_recommendation")
    def test_irrigation_recommendation_returns_direct_result(self, mock_get_irrigation_recommendation):
        mock_get_irrigation_recommendation.return_value = {
            "plan": {
                "frequencyPerWeek": 3,
                "durationMinutes": 25,
            },
            "raw_response": "{\"plan\": {\"frequencyPerWeek\": 3, \"durationMinutes\": 25}}",
        }

        response = self.client.post(
            "/api/rag/recommend/irrigation/",
            data={
                "farm_uuid": "sensor-123",
                "plant_name": "گوجه‌فرنگی",
                "growth_stage": "میوه‌دهی",
                "irrigation_method_name": "قطره‌ای",
            },
            format="json",
        )

        self.assertEqual(response.status_code, 200)
        self.assertEqual(response.json()["data"]["plan"]["frequencyPerWeek"], 3)
        mock_get_irrigation_recommendation.assert_called_once_with(
            farm_uuid="sensor-123",
            plant_name="گوجه‌فرنگی",
            growth_stage="میوه‌دهی",
            irrigation_method_name="قطره‌ای",
            query=None,
        )

    @patch("rag.services.fertilization.get_fertilization_recommendation")
    def test_fertilization_recommendation_returns_direct_result(self, mock_get_fertilization_recommendation):
        mock_get_fertilization_recommendation.return_value = {
            "plan": {
                "npkRatio": "20-20-20",
                "amountPerHectare": "80 kg",
            },
            "raw_response": "{\"plan\": {\"npkRatio\": \"20-20-20\", \"amountPerHectare\": \"80 kg\"}}",
        }

        response = self.client.post(
            "/api/rag/recommend/fertilization/",
            data={
                "farm_uuid": "sensor-456",
                "plant_name": "گندم",
                "growth_stage": "رویشی",
            },
            format="json",
        )

        self.assertEqual(response.status_code, 200)
        self.assertEqual(response.json()["data"]["plan"]["npkRatio"], "20-20-20")
        mock_get_fertilization_recommendation.assert_called_once_with(
            farm_uuid="sensor-456",
            plant_name="گندم",
            growth_stage="رویشی",
            query=None,
        )

    def test_removed_status_endpoints_return_404(self):
        irrigation_response = self.client.get("/api/rag/recommend/irrigation/sample-task/status/")
        fertilization_response = self.client.get("/api/rag/recommend/fertilization/sample-task/status/")

        self.assertEqual(irrigation_response.status_code, 404)
        self.assertEqual(fertilization_response.status_code, 404)