UPDATE

farm_ai_assistant/tests.py

@@ -2,6 +2,8 @@ from django.contrib.auth import get_user_model
 from django.test import TestCase, override_settings
 from rest_framework.test import APIRequestFactory, force_authenticate
 
+from farm_hub.models import FarmHub, FarmType
+
 from .models import Conversation, Message
 from .views import ChatTaskStatusView
 
@@ -16,24 +18,35 @@ class ChatTaskStatusViewTests(TestCase):
             email="farmer@example.com",
             phone_number="09120000000",
         )
+        self.farm_type, _ = FarmType.objects.get_or_create(name="زراعی")
+        self.farm = FarmHub.objects.create(
+            owner=self.user,
+            farm_type=self.farm_type,
+            name="Farm 1",
+        )
         self.conversation = Conversation.objects.create(
             owner=self.user,
+            farm=self.farm,
             title="Irrigation chat",
             farm_context={},
         )
         self.user_message = Message.objects.create(
             conversation=self.conversation,
+            farm=self.farm,
             role=Message.ROLE_USER,
             content="What is the best irrigation plan?",
             raw_response={
                 "task_id": "farm-ai-chat-task-123",
                 "status": "PENDING",
                 "status_url": "/api/tasks/farm-ai-chat-task-123/status/",
+                "farm_uuid": str(self.farm.farm_uuid),
             },
         )
 
     def test_status_success_uses_chat_mock_result_and_persists_assistant_message(self):
-        request = self.factory.get("/api/farm-ai-assistant/chat/task/farm-ai-chat-task-123/status/")
+        request = self.factory.get(
+            f"/api/farm-ai-assistant/chat/task/farm-ai-chat-task-123/status/?farm_uuid={self.farm.farm_uuid}"
+        )
         force_authenticate(request, user=self.user)
 
         response = ChatTaskStatusView.as_view()(request, task_id="farm-ai-chat-task-123")
@@ -43,6 +56,7 @@ class ChatTaskStatusViewTests(TestCase):
         self.assertEqual(response.data["data"]["task_id"], "farm-ai-chat-task-123")
         self.assertEqual(response.data["data"]["status"], "SUCCESS")
         self.assertEqual(response.data["data"]["conversation_id"], str(self.conversation.uuid))
+        self.assertEqual(response.data["data"]["farm_uuid"], str(self.farm.farm_uuid))
         self.assertEqual(response.data["data"]["result"]["content"], "Here is the recommended plan.")
         self.assertEqual(len(response.data["data"]["result"]["sections"]), 3)
         self.assertEqual(response.data["data"]["result"]["task_id"], "farm-ai-chat-task-123")
@@ -53,6 +67,8 @@ class ChatTaskStatusViewTests(TestCase):
             .first()
         )
         self.assertIsNotNone(assistant_message)
+        self.assertEqual(assistant_message.farm_id, self.farm.id)
         self.assertEqual(assistant_message.content, "Here is the recommended plan.")
         self.assertEqual(assistant_message.raw_response["task_id"], "farm-ai-chat-task-123")
+        self.assertEqual(assistant_message.raw_response["farm_uuid"], str(self.farm.farm_uuid))
         self.assertEqual(len(assistant_message.raw_response["sections"]), 3)