rq/tests/test_commands.py

import time
from multiprocessing import Process
from unittest import mock

from redis import Redis

from rq import Queue, Worker
from rq.command import send_command, send_kill_horse_command, send_shutdown_command, send_stop_job_command
from rq.exceptions import InvalidJobOperation, NoSuchJobError
from rq.serializers import JSONSerializer
from rq.worker import WorkerStatus
from tests import RQTestCase
from tests.fixtures import _send_kill_horse_command, _send_shutdown_command, long_running_job, raise_exc_mock


def start_work(queue_name, worker_name, connection_kwargs):
    worker = Worker(queue_name, name=worker_name, connection=Redis(**connection_kwargs))
    worker.work()


def start_work_burst(queue_name, worker_name, connection_kwargs):
    worker = Worker(queue_name, name=worker_name, connection=Redis(**connection_kwargs), serializer=JSONSerializer)
    worker.work(burst=True)


class TestCommands(RQTestCase):
    def test_shutdown_command(self):
        """Ensure that shutdown command works properly."""
        connection = self.connection
        worker = Worker('foo', connection=connection)

        p = Process(
            target=_send_shutdown_command, args=(worker.name, connection.connection_pool.connection_kwargs.copy())
        )
        p.start()
        worker.work()
        p.join(1)

    def test_pubsub_thread_survives_connection_error(self):
        """Ensure that the pubsub thread is still alive after its Redis connection is killed"""
        connection = self.connection
        worker = Worker('foo', connection=connection)
        worker.subscribe()

        assert worker.pubsub_thread.is_alive()

        for client in connection.client_list():
            connection.client_kill(client['addr'])

        time.sleep(0.0)  # Allow other threads to run
        assert worker.pubsub_thread.is_alive()

    def test_pubsub_thread_exits_other_error(self):
        """Ensure that the pubsub thread  exits on other than redis.exceptions.ConnectionError"""
        connection = self.connection
        worker = Worker('foo', connection=connection)

        with mock.patch('redis.client.PubSub.get_message', new_callable=raise_exc_mock):
            worker.subscribe()

        worker.pubsub_thread.join()
        assert not worker.pubsub_thread.is_alive()

    def test_kill_horse_command(self):
        """Ensure that shutdown command works properly."""
        connection = self.connection
        queue = Queue('foo', connection=connection)
        job = queue.enqueue(long_running_job, 4)
        worker = Worker('foo', connection=connection)

        p = Process(
            target=_send_kill_horse_command, args=(worker.name, connection.connection_pool.connection_kwargs.copy())
        )
        p.start()
        worker.work(burst=True)
        p.join(1)
        job.refresh()
        self.assertTrue(job.id in queue.failed_job_registry)

        p = Process(target=start_work, args=('foo', worker.name, connection.connection_pool.connection_kwargs.copy()))
        p.start()
        p.join(2)

        send_kill_horse_command(connection, worker.name)
        worker.refresh()
        # Since worker is not busy, command will be ignored
        self.assertEqual(worker.get_state(), WorkerStatus.IDLE)
        send_shutdown_command(connection, worker.name)

    def test_stop_job_command(self):
        """Ensure that stop_job command works properly."""

        connection = self.connection
        queue = Queue('foo', connection=connection, serializer=JSONSerializer)
        job = queue.enqueue(long_running_job, 3)
        worker = Worker('foo', connection=connection, serializer=JSONSerializer)

        # If job is not executing, an error is raised
        with self.assertRaises(InvalidJobOperation):
            send_stop_job_command(connection, job_id=job.id, serializer=JSONSerializer)

        # An exception is raised if job ID is invalid
        with self.assertRaises(NoSuchJobError):
            send_stop_job_command(connection, job_id='1', serializer=JSONSerializer)

        p = Process(
            target=start_work_burst, args=('foo', worker.name, connection.connection_pool.connection_kwargs.copy())
        )
        p.start()
        p.join(1)

        time.sleep(0.1)

        send_command(connection, worker.name, 'stop-job', job_id=1)
        time.sleep(0.25)
        # Worker still working due to job_id mismatch
        worker.refresh()
        self.assertEqual(worker.get_state(), WorkerStatus.BUSY)

        send_stop_job_command(connection, job_id=job.id, serializer=JSONSerializer)
        time.sleep(0.25)

        # Job status is set appropriately
        self.assertTrue(job.is_stopped)

        # Worker has stopped working
        worker.refresh()
        self.assertEqual(worker.get_state(), WorkerStatus.IDLE)