stop job test

Author: cunla

scheduler/tests/test_worker/test_worker_commands.py

@@ -1,9 +1,14 @@
+import json
+from threading import Thread
+from time import sleep
+
 from scheduler.helpers.queues import get_queue
-from scheduler.tests.jobs import test_job
+from scheduler.tests.jobs import test_job, two_seconds_job
+from ..test_job_decorator import long_running_func
 from ..test_views.base import BaseTestCase
-from ...redis_models import JobModel
+from ...redis_models import JobModel, JobStatus, WorkerModel
 from ...worker import create_worker
-from ...worker.commands import send_command
+from ...worker.commands import send_command, StopJobCommand
 from ...worker.commands.suspend_worker import SuspendWorkCommand
 
 
@@ -40,3 +45,29 @@ def test_stop_worker_command__bad_worker_name(self):
         self.assertFalse(worker._model.is_suspended)
         job = JobModel.get(job.name, connection=queue.connection)
         self.assertFalse(job.is_queued)
+
+    def test_stop_job_command__success(self):
+        # Arrange
+        worker_name = "test"
+        queue = get_queue("default")
+        job = queue.create_and_enqueue_job(two_seconds_job)
+        self.assertTrue(job.is_queued)
+        worker = create_worker("default", name=worker_name, burst=True, with_scheduler=False)
+        worker.bootstrap()
+
+        # Act
+        t = Thread(target=worker.work, args=(0,), name="worker-thread")
+        t.start()
+        sleep(0.1)
+        command = StopJobCommand(worker_name=worker_name, job_name=job.name)
+        command_payload = json.dumps(command.command_payload())
+        worker._command_listener.handle_payload(dict(data=command_payload))
+        worker.monitor_job_execution_process(job,queue)
+
+        # Assert
+        job = JobModel.get(job.name, connection=queue.connection)
+        worker = WorkerModel.get(worker.name, connection=queue.connection)
+        self.assertEqual(worker.stopped_job_name, job.name)
+        self.assertIsNone(worker.current_job_name)
+        self.assertEqual(job.status, JobStatus.STOPPED)
+        t.join()

scheduler/worker/commands/stop_job.py

@@ -32,12 +32,12 @@ def process_command(self, connection: ConnectionType) -> None:
         if job_model is None:
             logger.error(f"Job {self.job_name} not found")
             return
-        if worker_model.pid == worker_model.job_execution_process_pid:
-            logger.warning(f"Job execution process ID and worker process id {worker_model.pid} are equal, skipping")
-            return
         if not worker_model.job_execution_process_pid:
             logger.error(f"Worker {self.worker_name} has no job execution process")
             return
+        if worker_model.pid == worker_model.job_execution_process_pid:
+            logger.warning(f"Job execution process ID and worker process id {worker_model.pid} are equal, skipping")
+            return
         if worker_model.current_job_name != self.job_name:
             logger.info(
                 f"{self.worker_name} working on job {worker_model.current_job_name}, "

scheduler/worker/worker.py

@@ -581,7 +581,6 @@ def fork_job_execution_process(self, job: JobModel, queue: Queue) -> None:
             logger.debug(
                 f"[Worker {self.name}/{self._pid}]: Forking job execution process, job_execution_process_pid={child_pid}"
             )
-            refresh_queue_connection(queue)
             self._model.job_execution_process_pid = child_pid
             self._model.save(connection=queue.connection)
             self.procline(f"Forked {child_pid} at {time.time()}")
@@ -605,7 +604,6 @@ def monitor_job_execution_process(self, job: JobModel, queue: Queue) -> None:
         :param queue: The Queue
         """
         retpid = ret_val = None
-        job.started_at = utcnow()
         while True:
             try:
                 with SCHEDULER_CONFIG.DEATH_PENALTY_CLASS(