def test_shutdown_notifications_worker_shutdown_messages(raise_thread_exception, caplog): # capture all messages caplog.set_level(logging.NOTSET) # Given a middleware with two "threads" middleware = shutdown.ShutdownNotifications() middleware.manager.notifications = [1, 2] # Given a broker configured with the shutdown notifier broker = StubBroker(middleware=[middleware]) # When the worker is shutdown broker.emit_before("worker_shutdown", None) # Shutdown interrupts are raised in both threads raise_thread_exception.assert_has_calls([ mock.call(1, shutdown.Shutdown), mock.call(2, shutdown.Shutdown), ]) # And shutdown notifications are logged assert len(caplog.record_tuples) == 3 assert caplog.record_tuples == [ ("dramatiq.middleware.shutdown.ShutdownNotifications", logging.DEBUG, ( "Sending shutdown notification to worker threads..." )), ("dramatiq.middleware.shutdown.ShutdownNotifications", logging.INFO, ( "Worker shutdown notification. Raising exception in worker thread 1." )), ("dramatiq.middleware.shutdown.ShutdownNotifications", logging.INFO, ( "Worker shutdown notification. Raising exception in worker thread 2." )), ]
def test_shutdown_notifications_gevent_worker_shutdown_messages(caplog): # capture all messages caplog.set_level(logging.NOTSET) # Given a middleware with two threads middleware = shutdown.ShutdownNotifications() greenlet_1 = gevent.spawn() greenlet_2 = gevent.spawn() middleware.manager.notification_greenlets = [(1, greenlet_1), (2, greenlet_2)] # Given a broker configured with the shutdown notifier broker = StubBroker(middleware=[middleware]) # When the worker is shutdown broker.emit_before("worker_shutdown", None) # Shutdown interrupts are raised in both threads assert isinstance(greenlet_1.exception, shutdown.Shutdown) assert isinstance(greenlet_2.exception, shutdown.Shutdown) # And shutdown notifications are logged assert len(caplog.record_tuples) == 3 assert caplog.record_tuples == [ ("dramatiq.middleware.shutdown.ShutdownNotifications", logging.DEBUG, ( "Sending shutdown notification to worker threads..." )), ("dramatiq.middleware.shutdown.ShutdownNotifications", logging.INFO, ( "Worker shutdown notification. Raising exception in worker thread 1." )), ("dramatiq.middleware.shutdown.ShutdownNotifications", logging.INFO, ( "Worker shutdown notification. Raising exception in worker thread 2." )), ]
def test_dramatiq_completion( broker: StubBroker, worker: Worker, backend: WriterBackend, frozen_time: Any ) -> None: @dramatiq.actor(queue_name="test") def simple_task(a: str, b: str) -> str: return "hello" message = simple_task.send("a", b="b") assert backend.enqueued() == [ EnqueuedLog( type=LogType.ENQUEUED, timestamp=datetime.now(), job_id=message.message_id, task_id="simple_task", job=JobDetails( queue="test", task_path=( "tests.test_dramatiq.test_dramatiq_completion.<locals>" ".simple_task" ), execute_at=None, args=["a"], kwargs={"b": "b"}, options={}, ), ) ] assert backend.dequeued() == [] assert backend.completed() == [] worker.start() broker.join(simple_task.queue_name) worker.join() assert backend.dequeued() == [ DequeuedLog( job_id=message.message_id, task_id="simple_task", timestamp=datetime.now(), type=LogType.DEQUEUED, ) ] assert backend.completed() == [ CompletedLog( job_id=message.message_id, task_id="simple_task", timestamp=datetime.now(), result="hello", type=LogType.COMPLETED, ) ]
def stub_broker(): broker = StubBroker() broker.emit_after("process_boot") dramatiq.set_broker(broker) yield broker broker.flush_all() broker.close()
def test_shutdown_notifications_platform_not_supported(recwarn, monkeypatch): # monkeypatch fake platform to test logging. monkeypatch.setattr(shutdown, "current_platform", "not supported") # Given a broker configured with the shutdown notifier broker = StubBroker(middleware=[shutdown.ShutdownNotifications()]) # When the process boots broker.emit_after("process_boot") # A platform support warning is issued assert len(recwarn) == 1 assert str(recwarn[0].message) == ("ShutdownNotifications cannot kill threads " "on your current platform ('not supported').")
def test_time_limit_platform_not_supported(recwarn, monkeypatch): # monkeypatch fake platform to test logging. monkeypatch.setattr(time_limit, "current_platform", "not supported") # Given a broker configured with time limits broker = StubBroker(middleware=[time_limit.TimeLimit()]) # When the process boots broker.emit_after("process_boot") # A platform support warning is issued assert len(recwarn) == 1 assert str( recwarn[0].message) == ("TimeLimit cannot kill threads " "on your current platform ('not supported').")
def broker(sentry_init): sentry_init(integrations=[DramatiqIntegration()]) broker = StubBroker() broker.emit_after("process_boot") dramatiq.set_broker(broker) yield broker broker.flush_all() broker.close()
def test_create_range_for_stats_async( self, transactional_db, broker: stub.StubBroker, worker: dramatiq.Worker, hosting_provider_with_sample_user: ac_models.Hostingprovider, green_ip: gc_models.GreencheckIp, client, ): """ Create a collection of daily stats, for a range of dates provided. """ broker.declare_queue("default") generated_dates = self._set_up_dates_for_last_week() for date in generated_dates: gc = gc_factories.GreencheckFactory.create(date=date + relativedelta(hours=2)) # logger.info(f"gc {date}: {gc.__dict__}") logger.info(f"just this date: { generated_dates[0] }") gc_models.DailyStat.create_jobs_for_date_range_async( generated_dates, "total_count") # Wait for all the tasks to be processed broker.join("default") worker.join() green_stats = gc_models.DailyStat.objects.filter( green=gc_choices.BoolChoice.YES) grey_stats = gc_models.DailyStat.objects.filter( green=gc_choices.BoolChoice.NO) mixed_stats = gc_models.DailyStat.objects.exclude( green__in=[gc_choices.BoolChoice.YES, gc_choices.BoolChoice.NO]) # have we generated the expected stats per day? assert green_stats.count() == 7 assert grey_stats.count() == 7 assert mixed_stats.count() == 7 # we should one count showing zero green checks for each day assert [stat.count for stat in green_stats] == [0, 0, 0, 0, 0, 0, 0] # mixed and grey should be the same assert [stat.count for stat in grey_stats] == [1, 1, 1, 1, 1, 1, 1] assert [stat.count for stat in grey_stats] == [1, 1, 1, 1, 1, 1, 1]
def test_prometheus_middleware_exposes_metrics(): try: # Given a broker broker = StubBroker() # And an instance of the prometheus middleware prom = Prometheus() prom.after_process_boot(broker) # When I request metrics via HTTP with request.urlopen("http://127.0.0.1:9191") as resp: # Then the response should be successful assert resp.getcode() == 200 finally: prom.after_worker_shutdown(broker, None)
def get_dramatiq_broker_object(self): """ This method initializes the broker object for Dramatiq and saves it in Django's settings. """ if self.DRAMATIQ_BROKER_URL: url = self.DRAMATIQ_BROKER_URL.strip() if url is None: kind, host, port = "stub", None, None else: kind, _, url = url.partition("://") host, _, port = url.partition(":") host = host or None port = int(port) if port else None else: kind = self.DRAMATIQ_BROKER_TYPE host = self.DRAMATIQ_BROKER_HOST or None port = self.DRAMATIQ_BROKER_PORT or None # Separate non-null args kwargs = [("host", host), ("port", port)] kwargs = {k: v for k, v in kwargs if v is not None} # Initializes broker if kind == "stub": from dramatiq.brokers.stub import StubBroker broker = StubBroker() elif kind == "redis": from dramatiq.brokers.redis import RedisBroker broker = RedisBroker(**kwargs) elif kind == "rabbitmq": from dramatiq.brokers.rabbitmq import RabbitmqBroker broker = RabbitmqBroker(**kwargs) else: raise ValueError(f"invalid dramatiq broker: {kind}") # Configure as default and exit dramatiq.set_broker(broker) return broker
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 test_create_stat_async( self, transactional_db, broker: stub.StubBroker, worker: dramatiq.Worker, hosting_provider_with_sample_user: ac_models.Hostingprovider, green_ip: gc_models.GreencheckIp, client, ): """ Create a collection of daily stats, for a range of dates provided, but have a worker create the stats asynchronously. """ broker.declare_queue("default") assert gc_models.DailyStat.objects.count() == 0 # set up our date range generated_dates = self._set_up_dates_for_last_week() for date in generated_dates: gc_factories.GreencheckFactory.create(date=date + relativedelta(hours=2)) chosen_date = str(generated_dates[0].date()) # we use the 'send' with the 'transactional_db' fixture here instead of db # because if we use the regular db fixture, the workers can not see what is # happening 'inside' this test. TODO: check that this really is the # explanation for this strange test behaviour gc_tasks.create_stat_async.send(date_string=chosen_date, query_name="total_count") # import ipdb # ipdb.set_trace() # Wait for all the tasks to be processed broker.join(gc_tasks.create_stat_async.queue_name) worker.join() # import ipdb # ipdb.set_trace() # hae we generate the daily stats? assert gc_models.DailyStat.objects.count() == 3 # do that they have the right date? for stat in gc_models.DailyStat.objects.all(): assert str(stat.stat_date) == chosen_date # do the stats count up to what we expect? green_daily_stat = gc_models.DailyStat.objects.filter( green=gc_choices.BoolChoice.YES).first() grey_daily_stat = gc_models.DailyStat.objects.filter( green=gc_choices.BoolChoice.NO).first() mixed_daily_stat = gc_models.DailyStat.objects.exclude( green__in=[gc_choices.BoolChoice.YES, gc_choices.BoolChoice.NO ]).first() assert green_daily_stat.count == 0 assert grey_daily_stat.count == 1 assert mixed_daily_stat.count == 1
def broker(backend: WriterBackend) -> StubBroker: stub_broker = StubBroker(middleware=[TaskLogsMiddleware(backend=backend)]) stub_broker.emit_after("process_boot") dramatiq.set_broker(stub_broker) return stub_broker
class InitDB(dramatiq.Middleware): def before_worker_boot(self, broker, worker): async def run(): from . import db await db.db.set_bind(db.CONNECTION_STR) loop.run_until_complete(run()) def before_worker_shutdown(self, broker, worker): shutdown.set() MIDDLEWARE = [ m() for m in (AgeLimit, ShutdownNotifications, Callbacks, Pipelines, Retries) ] if TEST: broker = StubBroker(middleware=MIDDLEWARE) broker.emit_after("process_boot") else: broker = RedisBroker( connection_pool=redis.ConnectionPool.from_url(REDIS_HOST), middleware=MIDDLEWARE) broker.add_middleware(InitDB()) dramatiq.set_broker(broker)
import os import signal import time from dramatiq.brokers.stub import StubBroker from .common import skip_on_windows broker = StubBroker() def remove(filename): try: os.remove(filename) except OSError: pass @skip_on_windows def test_cli_scrubs_stale_pid_files(start_cli): try: # Given that I have an existing file containing an old pid filename = "test_scrub.pid" with open(filename, "w") as f: f.write("999999") # When I try to start the cli and pass that file as a PID file proc = start_cli("tests.test_pidfile:broker", extra_args=["--pid-file", filename]) # And I wait for it to write the pid file time.sleep(1)
import os from dramatiq.brokers.rabbitmq import RabbitmqBroker from dramatiq.brokers.stub import StubBroker if os.getenv("UNIT_TESTS") == "1": broker = StubBroker(worker_timeout=1000, worker_threads=1) broker.emit_after("process_boot") else: broker = RabbitmqBroker()
from lib.telegram import TelegramBot from lib.telegram.command import Command logger = logging.getLogger(__name__) logger.setLevel(settings.LOG_LEVEL) rate_limits_backend: RateLimiterBackend broker: dramatiq.Broker result_backend: ResultBackend if settings.UNIT_TESTS: # Setup backends result_backend = ResultStubBackend() rate_limits_backend = RateLimitsStubBackend() # Setup brokers broker = StubBroker() else: logger.info("Setup redis broker") # Setup backends result_backend = ResultRedisBackend(url=settings.RESULT_REDIS_URL) rate_limits_backend = RateLimitsRedisBackend( url=settings.RATE_LIMITS_REDIS_URL) # Setup brokers broker = RedisBroker(url=settings.BROKER_REDIS_URL) results = Results(backend=result_backend) broker.add_middleware(results) if settings.PROMETHEUS_METRICS_SERVER_ENABLE: broker.add_middleware( Prometheus(
def stub_broker(): from dramatiq.brokers.stub import StubBroker broker = StubBroker() actors.initialize_broker(broker) yield broker
import contextlib from unittest.mock import MagicMock import dramatiq import pytest from dramatiq import Worker from dramatiq.brokers.stub import StubBroker from sqlalchemy import MetaData, Table from events_api import tasks from events_api.app import create_app from events_api.core.project import database broker = StubBroker() broker.emit_after("process_boot") def clear_tables(engine, table_list=None): meta = MetaData() with contextlib.closing(engine.connect()) as con: trans = con.begin() if table_list is None: meta.reflect(bind=engine) for table in reversed(meta.sorted_tables): con.execute(table.delete()) else: for table_name in table_list: table = Table(table_name, meta, autoload=True, autoload_with=engine)
from pathlib import PureWindowsPath, PurePosixPath import docker as _docker import dramatiq import socketio from dramatiq.brokers.rabbitmq import RabbitmqBroker from dramatiq.brokers.stub import StubBroker from api import config docker = _docker.from_env() broker = StubBroker() if config.UNIT_TESTING else RabbitmqBroker( host=config.RABBITMQ_HOST) dramatiq.set_broker(broker) external_sio = socketio.AsyncAioPikaManager(write_only=True) HOST_PATH_TYPE = PureWindowsPath if config.DOCKER_HOST_OS.lower( ) == 'windows' else PurePosixPath