def test_abort_before_processing(
stub_broker: dramatiq.Broker,
stub_event_backend: EventBackend,
mode: AbortMode,
) -> None:
calls = []
abortable = Abortable(backend=stub_event_backend)
stub_broker.add_middleware(abortable)
@dramatiq.actor(abortable=True, max_retries=0)
def do_work() -> None:
calls.append(1)
stub_broker.emit_after("process_boot")
# If I send it a message
message = do_work.send()
# And cancel right after.
abort(message.message_id, mode=mode)
# Then start the worker.
worker = dramatiq.Worker(stub_broker, worker_timeout=100, worker_threads=1)
worker.start()
stub_broker.join(do_work.queue_name)
worker.join()
worker.stop()
# I expect the task to not have been called.
assert sum(calls) == 0
def worker(broker: Broker) -> Generator[dramatiq.Worker, None, None]:
worker = dramatiq.Worker(broker, worker_timeout=100)
worker.start()
yield worker
worker.stop()
def worker(broker):
worker = dramatiq.Worker(broker, worker_threads=THREAD_COUNT)
worker.start()
try:
yield worker
finally:
worker.stop()
def test_add_middleware(app, broker, run_mock):
class TestMiddleware(Middleware):
@property
def actor_options(self):
return {"test_option"}
def after_process_message(self,
broker,
message,
*,
result=None,
exception=None):
run_mock()
broker.add_middleware(TestMiddleware())
@broker.actor(test_option=123)
def task():
pass
broker.init_app(app)
broker.add_middleware(TestMiddleware())
assert task.options['test_option'] == 123
task.send()
worker = dramatiq.Worker(broker)
worker.start()
worker.join()
assert run_mock.call_count == 2
def stub_worker():
broker = dramatiq.get_broker()
worker = dramatiq.Worker(broker, worker_timeout=1000)
worker.start()
try:
yield worker
finally:
worker.stop()
def test_dramatiq_task(app, broker, dramatiq_task, run_mock):
with pytest.raises(RuntimeError, match=r'init_app\(\) must be called'):
dramatiq_task.send()
broker.init_app(app)
dramatiq_task.send()
worker = dramatiq.Worker(broker)
worker.start()
worker.join()
run_mock.assert_called_once()
def test_immediate_init(app, run_mock):
broker = StubBroker(app, config_prefix='IMMEDIATE_INIT_BROKER')
@broker.actor
def task():
run_mock()
task.send()
worker = dramatiq.Worker(broker)
worker.start()
worker.join()
run_mock.assert_called_once()
def test_register_task_after_init(app, broker, run_mock):
broker.init_app(app)
@broker.actor
def task(p1, p2):
run_mock(p1, p2)
task.send('param1', 'param2')
worker = dramatiq.Worker(broker)
worker.start()
worker.join()
run_mock.assert_called_once_with('param1', 'param2')
def test_lazy_actor(app, run_mock):
import dramatiq.brokers.stub
broker = dramatiq.brokers.stub.StubBroker()
@dramatiq.actor(actor_class=LazyActor, broker=broker)
def task():
run_mock()
task.send()
worker = dramatiq.Worker(broker)
worker.start()
worker.join()
run_mock.assert_called_once()
def main(args):
broker = dramatiq.get_broker()
broker.emit_after("process_boot")
worker = dramatiq.Worker(broker, worker_threads=1)
worker.start()
while True:
try:
time.sleep(1)
except KeyboardInterrupt:
break
worker.stop()
broker.close()
def test_flask_app_context(app, broker, run_mock):
@broker.actor
def task():
assert app.config is flask.current_app.config
run_mock()
broker.init_app(app)
assert [
1 for m in broker.middleware
if type(m).__name__ == 'AppContextMiddleware'
]
task.send()
worker = dramatiq.Worker(broker)
worker.start()
worker.join()
run_mock.assert_called_once()
def setUp(self):
super().setUp()
self.TEST_USER_NAME = 'testuser'
self.TEST_USER_PASSWORD = '******'
# Create a user
test_user1 = User.objects.create_user(
username=self.TEST_USER_NAME,
password=self.TEST_USER_PASSWORD,
)
test_user1.save()
print(f'TEST USER {test_user1.username}')
print(f'TEST USER {test_user1.id}')
self.test_user = test_user1
self.broker = dramatiq.get_broker()
self.worker = dramatiq.Worker(self.broker, worker_timeout=100)
self.worker.start()
def test_generic_actor(app, broker, run_mock):
class AbstractTask(dramatiq.GenericActor):
class Meta:
abstract = True
def perform(self, arg):
self.run(arg)
class Task(AbstractTask):
def run(self, arg):
run_mock(arg)
broker.init_app(app)
Task.send('message')
worker = dramatiq.Worker(broker)
worker.start()
worker.join()
run_mock.assert_called_once_with('message')
def stub_worker(monkeypatch, settings, _runner):
if settings.JOEFLOW_TASK_RUNNER == "joeflow.runner.celery.task_runner":
yield mock.Mock()
else:
import dramatiq
broker = dramatiq.get_broker()
broker.emit_after("process_boot")
broker.flush_all()
worker = dramatiq.Worker(broker, worker_timeout=100)
worker.start()
class Meta:
@staticmethod
def wait():
broker.join(settings.JOEFLOW_CELERY_QUEUE_NAME, timeout=60000)
worker.join()
yield Meta
worker.stop()
def app_with_scout(config=None):
"""
Context manager that configures a Dramatiq app with Scout middleware
installed.
"""
# Enable Scout by default in tests.
if config is None:
config = {"monitor": True}
# Disable running the agent.
config["core_agent_launch"] = False
broker = StubBroker()
broker.emit_after("process_boot")
dramatiq.set_broker(broker)
@dramatiq.actor(max_retries=0)
def hello():
return "Hello World!"
@dramatiq.actor(max_retries=0)
def fail():
raise ValueError("BØØM!") # non-ASCII
worker = dramatiq.Worker(broker, worker_timeout=0)
# Setup according to https://docs.scoutapm.com/#dramatiq
Config.set(**config)
broker.add_middleware(ScoutMiddleware(),
before=broker.middleware[0].__class__)
worker.start()
App = namedtuple("App", ["broker", "worker", "hello", "fail"])
try:
yield App(broker=broker, worker=worker, hello=hello, fail=fail)
finally:
worker.stop()
# Reset Scout configuration.
Config.reset_all()
def worker(broker):
worker = dramatiq.Worker(broker, worker_timeout=100)
worker.start()
yield worker
worker.stop()
def worker(broker):
worker = dramatiq.Worker(broker)
worker.start()
yield worker
worker.stop()
def setUp(self):
super().setUp()
self.broker = dramatiq.get_broker()
self.worker = dramatiq.Worker(self.broker, worker_timeout=100)
self.worker.start()
def process_messages():
worker = dramatiq.Worker(redis_broker)
worker.start()
redis_broker.join(throughput.queue_name)
worker.stop()
def process_messages():
worker = dramatiq.Worker(rabbitmq_broker)
worker.start()
rabbitmq_broker.join(rabbitmq_random_queue)
worker.stop()
