From 500fa58cbeb58a770a9f57ae483f1ed07d7c2bdd Mon Sep 17 00:00:00 2001
From: shamoon <4887959+shamoon@users.noreply.github.com>
Date: Tue, 25 Feb 2025 11:28:34 -0800
Subject: [PATCH] Mock this way

---
 src/documents/tests/test_api_tasks.py | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/src/documents/tests/test_api_tasks.py b/src/documents/tests/test_api_tasks.py
index 57b56a2ef..c139d05da 100644
--- a/src/documents/tests/test_api_tasks.py
+++ b/src/documents/tests/test_api_tasks.py
@@ -9,6 +9,7 @@ from rest_framework.test import APITestCase
 
 from documents.models import PaperlessTask
 from documents.tests.utils import DirectoriesMixin
+from documents.views import TasksViewSet
 
 
 class TestTasks(DirectoriesMixin, APITestCase):
@@ -311,8 +312,7 @@ class TestTasks(DirectoriesMixin, APITestCase):
 
         self.assertEqual(returned_data["related_document"], "1234")
 
-    @mock.patch("documents.tasks.train_classifier")
-    def test_run_train_classifier_task(self, mock_train_classifier):
+    def test_run_train_classifier_task(self):
         """
         GIVEN:
             - A superuser
@@ -321,7 +321,13 @@ class TestTasks(DirectoriesMixin, APITestCase):
         THEN:
             - The task is run
         """
-        mock_train_classifier.return_value = "Task started"
+        mock_train_classifier = mock.Mock(return_value="Task started")
+        TasksViewSet.TASK_AND_ARGS_BY_NAME = {
+            PaperlessTask.TaskName.TRAIN_CLASSIFIER: (
+                mock_train_classifier,
+                {"scheduled": False},
+            ),
+        }
         response = self.client.post(
             self.ENDPOINT + "run/",
             {"task_name": PaperlessTask.TaskName.TRAIN_CLASSIFIER},