mirror of https://github.com/rq/rq.git
129 lines
4.7 KiB
Python
129 lines
4.7 KiB
Python
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)
|