def get_mware():
"""Returns a broker middleware designed to support making available actor execution results to downstream actors.
"""
return Results(
RedisBackend(client=redis.Redis(
db=EnvVars.DB, host=EnvVars.HOST, port=EnvVars.PORT)))
def test_actor_no_warning_when_returns_result_while_piping_and_store_results_is_not_set(
stub_broker, stub_worker):
# Given a result backend
# And a broker with the results middleware
stub_broker.add_middleware(Results(backend=StubBackend()))
# And that I've mocked the logging class
with patch("logging.Logger.warning") as warning_mock:
# And I have an actor that always returns 1, and does not store results
@dramatiq.actor
def always_1():
return 1
# And an actor that takes a single argument
@dramatiq.actor
def noop(x):
pass
# When I send that actor a message
(always_1.message() | noop.message()).run()
# And wait for the message to get processed
stub_broker.join(always_1.queue_name)
stub_worker.join()
# Then a warning should not be logged
warning_messages = [args[0] for _, args, _ in warning_mock.mock_calls]
assert not any("the value has been discarded" in x
for x in warning_messages)
def stub_broker(r):
result_backend = RedisBackend()
result_backend.client = r
broker.add_middleware(Results(backend=result_backend))
broker.flush_all()
dramatiq.set_broker(stub_broker)
return broker
def configure():
result_backend = RedisBackend(port=63791)
rabbitmq_broker = RabbitmqBroker(port=56721)
rabbitmq_broker.add_middleware(Results(backend=result_backend))
dramatiq.set_broker(rabbitmq_broker)
def test_pipe_ignore_applies_to_receiving_message(stub_broker, stub_worker, result_backend):
# Given a result backend
# And a broker with the results middleware
stub_broker.add_middleware(Results(backend=result_backend))
@dramatiq.actor(store_results=True)
def return_args(*args):
return args
# When I compose pipe of three messages with pipe_ignore option on second message
pipe = (
return_args.message(1) |
return_args.message_with_options(pipe_ignore=True, args=(2, )) |
return_args.message(3)
)
# And then run and wait for it to complete
pipe.run()
stub_broker.join(return_args.queue_name)
results = list(pipe.get_results())
# The then result of the first message should NOT be passed as
# argument to the second message and the result of the second
# message should be passed as argument to the third message.
assert results == [[1], [2], [3, [2]]]
def test_groups_execute_inner_groups(stub_broker, stub_worker, backend, result_backends):
# Given that I have a result backend
backend = result_backends[backend]
stub_broker.add_middleware(Results(backend=backend))
# And I have an actor that sleeps for 100ms
@dramatiq.actor(store_results=True)
def wait():
time.sleep(0.1)
# When I group multiple groups inside one group and run it
t = time.monotonic()
g = group(group(wait.message() for _ in range(2)) for _ in range(3))
g.run()
# And wait on the group to complete
results = list(g.get_results(block=True))
# Then the total elapsed time should be less than 500ms
assert time.monotonic() - t <= 0.5
# And I should get back 3 results each with 2 results inside it
assert results == [[None, None]] * 3
# And the group should be completed
assert g.completed
def test_groups_execute_jobs_in_parallel(stub_broker, stub_worker, backend, result_backends):
# Given that I have a result backend
backend = result_backends[backend]
stub_broker.add_middleware(Results(backend=backend))
# And I have an actor that sleeps for 100ms
@dramatiq.actor(store_results=True)
def wait():
time.sleep(0.1)
# When I group multiple of these actors together and run them
t = time.monotonic()
g = group([wait.message() for _ in range(5)])
g.run()
# And wait on the group to complete
results = list(g.get_results(block=True))
# Then the total elapsed time should be less than 500ms
assert time.monotonic() - t <= 0.5
# And I should get back as many results as there were jobs in the group
assert len(results) == len(g)
# And the group should be completed
assert g.completed
def test_groups_expose_completion_stats(stub_broker, stub_worker, backend, result_backends):
# Given that I have a result backend
backend = result_backends[backend]
stub_broker.add_middleware(Results(backend=backend))
# And an actor that waits some amount of time
condition = Condition()
@dramatiq.actor(store_results=True)
def wait(n):
time.sleep(n)
with condition:
condition.notify_all()
return n
# When I group messages of varying durations together and run the group
g = group(wait.message(n) for n in range(1, 4))
g.run()
# Then every time a job in the group completes, the completed_count should increase
for count in range(1, len(g) + 1):
with condition:
condition.wait(5)
time.sleep(0.1) # give the worker time to set the result
assert g.completed_count == count
# Finally, completed should be true
assert g.completed
def test_pipelines_expose_completion_stats(stub_broker, stub_worker, backend, result_backends):
# Given a result backend
backend = result_backends[backend]
# And a broker with the results middleware
stub_broker.add_middleware(Results(backend=backend))
# And an actor that waits some amount of time
condition = Condition()
@dramatiq.actor(store_results=True)
def wait(n):
time.sleep(n)
with condition:
condition.notify_all()
return n
# When I pipe some messages intended for that actor together and run the pipeline
pipe = wait.message(1) | wait.message()
pipe.run()
# Then every time a job in the pipeline completes, the completed_count should increase
for count in range(1, len(pipe) + 1):
with condition:
condition.wait(2)
time.sleep(0.1) # give the worker time to set the result
assert pipe.completed_count == count
# Finally, completed should be true
assert pipe.completed
def ready(self):
dramatiq.set_encoder(self.select_encoder())
result_backend_settings = self.result_backend_settings()
if result_backend_settings:
result_backend_path = result_backend_settings.get(
"BACKEND", "dramatiq.results.backends.StubBackend")
result_backend_class = load_class(result_backend_path)
result_backend_options = result_backend_settings.get(
"BACKEND_OPTIONS", {})
result_backend = result_backend_class(**result_backend_options)
results_middleware_options = result_backend_settings.get(
"MIDDLEWARE_OPTIONS", {})
results_middleware = Results(backend=result_backend,
**results_middleware_options)
else:
result_backend = None
results_middleware = None
broker_settings = self.broker_settings()
broker_path = broker_settings["BROKER"]
broker_class = load_class(broker_path)
broker_options = broker_settings.get("OPTIONS", {})
middleware = [
load_middleware(path)
for path in broker_settings.get("MIDDLEWARE", [])
]
if result_backend is not None:
middleware.append(results_middleware)
broker = broker_class(middleware=middleware, **broker_options)
dramatiq.set_broker(broker)
def test_messages_without_actor_not_crashing_lookup_options(
stub_broker, redis_result_backend):
message = Message(
queue_name="default",
actor_name="idontexist",
args=(),
kwargs={},
options={},
)
assert Results(backend=redis_result_backend).after_nack(
stub_broker, message) is None
def redis_broker(r):
broker = RedisBroker()
broker.client = r
broker.client.flushall()
result_backend = RedisBackend()
result_backend.client = r
broker.add_middleware(Results(backend=result_backend))
broker.emit_after("process_boot")
# monkeypatch.setattr("yatsm.jobs.redis_broker", broker)
dramatiq.set_broker(broker)
yield broker
broker.client.flushall()
broker.close()
def setup_redis_broker() -> None:
_connection_pool: redis.BlockingConnectionPool = redis.BlockingConnectionPool(
host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB, password=REDIS_PASSWORD
)
_redis_db: redis.StrictRedis = redis.StrictRedis(connection_pool=_connection_pool)
_result_backend = RedisBackend(encoder=PickleEncoder(), client=_redis_db)
_result_middleware = Results(backend=_result_backend)
broker: Broker = RedisBroker(
connection_pool=_connection_pool,
middleware=[_result_middleware],
namespace="lrw",
)
dramatiq.set_broker(broker)
dramatiq.set_encoder(dramatiq.PickleEncoder())
def initialize(cls):
global RATE_LIMITER_BACKEND
dramatiq.set_encoder(cls.select_encoder())
result_backend_settings = cls.result_backend_settings()
if result_backend_settings:
result_backend_path = result_backend_settings.get(
"BACKEND", "dramatiq.results.backends.StubBackend")
result_backend_class = import_string(result_backend_path)
result_backend_options = result_backend_settings.get(
"BACKEND_OPTIONS", {})
result_backend = result_backend_class(**result_backend_options)
results_middleware_options = result_backend_settings.get(
"MIDDLEWARE_OPTIONS", {})
results_middleware = Results(backend=result_backend,
**results_middleware_options)
else:
result_backend = None
results_middleware = None
rate_limiter_backend_settings = cls.rate_limiter_backend_settings()
if rate_limiter_backend_settings:
rate_limiter_backend_path = rate_limiter_backend_settings.get(
"BACKEND", "dramatiq.rate_limits.backends.stub.StubBackend")
rate_limiter_backend_class = import_string(
rate_limiter_backend_path)
rate_limiter_backend_options = rate_limiter_backend_settings.get(
"BACKEND_OPTIONS", {})
RATE_LIMITER_BACKEND = rate_limiter_backend_class(
**rate_limiter_backend_options)
broker_settings = cls.broker_settings()
broker_path = broker_settings["BROKER"]
broker_class = import_string(broker_path)
broker_options = broker_settings.get("OPTIONS", {})
middleware = [
load_middleware(path)
for path in broker_settings.get("MIDDLEWARE", [])
]
if result_backend is not None:
middleware.append(results_middleware)
broker = broker_class(middleware=middleware, **broker_options)
dramatiq.set_broker(broker)
def test_pipelines_can_be_incomplete(stub_broker, result_backend):
# Given that I am not running a worker
# And I have a result backend
stub_broker.add_middleware(Results(backend=result_backend))
# And I have an actor that does nothing
@dramatiq.actor(store_results=True)
def do_nothing():
return None
# And I've run a pipeline
pipe = do_nothing.message() | do_nothing.message_with_options(pipe_ignore=True)
pipe.run()
# When I check if the pipeline has completed
# Then it should return False
assert not pipe.completed
def test_messages_can_get_results_from_inferred_backend(
stub_broker, stub_worker, result_backend):
# Given a result backend
# And a broker with the results middleware
stub_broker.add_middleware(Results(backend=result_backend))
# And an actor that stores a result
@dramatiq.actor(store_results=True)
def do_work():
return 42
# When I send that actor a message
message = do_work.send()
# And wait for a result
# Then I should get that result back
assert message.get_result(block=True) == 42
def test_actors_can_store_results(stub_broker, stub_worker, result_backend):
# Given a result backend
# And a broker with the results middleware
stub_broker.add_middleware(Results(backend=result_backend))
# And an actor that stores results
@dramatiq.actor(store_results=True)
def do_work():
return 42
# When I send that actor a message
message = do_work.send()
# And wait for a result
result = result_backend.get_result(message, block=True)
# Then the result should be what the actor returned
assert result == 42
def test_messages_can_fail_and_propagate(stub_broker, stub_worker,
result_backend):
# Given a result backend
# And a broker with the results middleware
stub_broker.add_middleware(Results(backend=result_backend))
# And an actor that stores a result
@dramatiq.actor(store_results=True)
def do_work():
raise TestActorException('msg')
# When I send that actor a message
message = do_work.send()
# And wait for a result
# Then I should get that result back
with pytest.raises(TestActorException, match='msg'):
message.get_result(block=True)
def test_retrieving_a_result_can_time_out(stub_broker, stub_worker,
result_backend):
# Given a result backend
# And a broker with the results middleware
stub_broker.add_middleware(Results(backend=result_backend))
# And an actor that sleeps for a long time before it stores a result
@dramatiq.actor(store_results=True)
def do_work():
time.sleep(0.2)
return 42
# When I send that actor a message
message = do_work.send()
# And wait for a result
# Then a ResultTimeout error should be raised
with pytest.raises(ResultTimeout):
result_backend.get_result(message, block=True, timeout=100)
def test_pipeline_respects_own_delay(stub_broker, stub_worker, result_backend):
# Given a result backend
# And a broker with the results middleware
stub_broker.add_middleware(Results(backend=result_backend))
# And an actor that adds two numbers together
@dramatiq.actor(store_results=True)
def add(x, y):
return x + y
# When I pipe some messages intended for that actor together and run the pipeline with a delay
pipe = add.message(1, 2) | add.message(3)
pipe.run(delay=10000)
# And get the results with a lower timeout than the the pipeline is delayed by
# Then a ResultTimeout error should be raised
with pytest.raises(ResultTimeout):
for _ in pipe.get_results(block=True, timeout=100):
pass
def test_pipeline_results_can_be_retrieved(stub_broker, stub_worker, result_backend):
# Given a result backend
# And a broker with the results middleware
stub_broker.add_middleware(Results(backend=result_backend))
# And an actor that adds two numbers together and stores the result
@dramatiq.actor(store_results=True)
def add(x, y):
return x + y
# When I pipe some messages intended for that actor together and run the pipeline
pipe = add.message(1, 2) | (add.message(3) | add.message(4))
pipe.run()
# Then the pipeline result should be the sum of 1, 2, 3 and 4
assert pipe.get_result(block=True) == 10
# And I should be able to retrieve individual results
assert list(pipe.get_results()) == [3, 6, 10]
def test_actors_can_store_exceptions(stub_broker, stub_worker, result_backend):
# Given a result backend
# And a broker with the results middleware
stub_broker.add_middleware(Results(backend=result_backend))
# And an actor that stores results
@dramatiq.actor(store_results=True, max_retries=0)
def do_work():
raise RuntimeError("failed")
# When I send that actor a message
message = do_work.send()
# And wait for a result
# Then the result should be an exception
with pytest.raises(ResultFailure) as e:
result_backend.get_result(message, block=True)
assert str(e.value) == "actor raised RuntimeError: failed"
def test_groups_can_time_out(stub_broker, stub_worker, result_backend):
# Given that I have a result backend
stub_broker.add_middleware(Results(backend=result_backend))
# And I have an actor that sleeps for 300ms
@dramatiq.actor(store_results=True)
def wait():
time.sleep(0.3)
# When I group a few jobs together and run it
g = group(wait.message() for _ in range(2))
g.run()
# And wait for the group to complete with a timeout
# Then a ResultTimeout error should be raised
with pytest.raises(ResultTimeout):
g.wait(timeout=100)
# And the group should not be completed
assert not g.completed
def test_pipeline_results_respect_timeouts(stub_broker, stub_worker, result_backend):
# Given a result backend
# And a broker with the results middleware
stub_broker.add_middleware(Results(backend=result_backend))
# And an actor that waits some amount of time then doubles that amount
@dramatiq.actor(store_results=True)
def wait(n):
time.sleep(n)
return n * 2
# When I pipe some messages intended for that actor together and run the pipeline
pipe = wait.message(1) | wait.message() | wait.message()
pipe.run()
# And get the results with a lower timeout than the tasks can complete in
# Then a ResultTimeout error should be raised
with pytest.raises(ResultTimeout):
for _ in pipe.get_results(block=True, timeout=1000):
pass
def test_actors_results_are_backwards_compatible(stub_broker, stub_worker,
result_backend):
# Given a result backend
# And a broker with the results middleware
stub_broker.add_middleware(Results(backend=result_backend))
# And an actor that stores results
@dramatiq.actor(store_results=True)
def do_work():
return 42
# And I have a result created using an old version of dramatiq
message = do_work.message()
message_key = result_backend.build_message_key(message)
result_backend._store(message_key, 42, 3600000)
# When I grab that result
result = result_backend.get_result(message, block=True)
# Then it should be unwrapped correctly
assert result == 42
def test_pipeline_respects_bigger_of_first_messages_and_pipelines_delay(
stub_broker, stub_worker, result_backend):
# Given a result backend
# And a broker with the results middleware
stub_broker.add_middleware(Results(backend=result_backend))
# And an actor that adds two numbers together
@dramatiq.actor(store_results=True)
def add(x, y):
return x + y
# When I pipe some messages intended for that actor together, first of which is delayed
# And the pipeline is delayed with a bigger value than the first message, and run the pipeline
pipe = add.message_with_options(args=(1, 2), delay=100) | add.message(3)
pipe.run(delay=10000)
# And get the results with a higher timeout than first message's delay, but lower than pipeline's delay
# Then a ResultTimeout error should be raised
with pytest.raises(ResultTimeout):
for _ in pipe.get_results(block=True, timeout=300):
pass
def test_retrieving_a_result_can_return_not_ready(stub_broker, stub_worker,
backend, result_backends):
# Given a result backend
backend = result_backends[backend]
# And a broker with the results middleware
stub_broker.add_middleware(Results(backend=backend))
# And an actor that sleeps for a long time before it stores a result
@dramatiq.actor(store_results=True)
def do_work():
time.sleep(0.2)
return 42
# When I send that actor a message
message = do_work.send()
# And get the result without blocking
# Then a ResultMissing error should be raised
with pytest.raises(ResultMissing):
backend.get_result(message)
def setup_dramatiq(config_class=Config):
global broker, results_backend, rate_limiter_backend
o = urlsplit(config_class.REDIS_URL)
results_backend = RedisResultsBackend(
host=o.hostname,
port=o.port,
db=o.path.split('/')[-1] or None,
password=o.password
)
rate_limiter_backend = RedisRateLimiterBackend(
host=o.hostname,
port=o.port,
db=o.path.split('/')[-1] or None,
password=o.password
)
broker = RedisBroker(url=config_class.REDIS_URL)
broker.add_middleware(Results(backend=results_backend))
dramatiq.set_broker(broker)
def test_actor_no_warning_when_returns_result_and_results_middleware_present(
stub_broker, stub_worker, result_backend):
# Given a result backend
# And a broker with the results middleware
stub_broker.add_middleware(Results(backend=result_backend))
# And that I've mocked the logging class
with patch("logging.Logger.warning") as warning_mock:
# And I have an actor that always returns 1, and does store results
@dramatiq.actor(store_results=True)
def always_1():
return 1
# When I send that actor a message
always_1.send()
# And wait for the message to get processed
stub_broker.join(always_1.queue_name)
stub_worker.join()
# Then a warning should not be logged
warning_messages = [args[0] for _, args, _ in warning_mock.mock_calls]
assert not any("Consider adding the Results middleware" in x
for x in warning_messages)
def setup_broker_and_backend():
redis_broker = RedisBroker(host="localhost", port=6379)
result_backend = RedisBackend()
dramatiq.set_broker(redis_broker)
redis_broker.add_middleware(Results(backend=result_backend))
return redis_broker
