import logging import pytest from taskiq import InMemoryBroker, ScheduleSource, TaskiqScheduler from taskiq.exceptions import ScheduledTaskCancelledError from taskiq_deduplication import RedisDeduplicationMiddleware from taskiq_deduplication.middleware import ( DEDUP_EXPLICIT_KEY_LABEL, DEDUP_KEY_FIELDS_LABEL, DEDUP_LABEL, ) from taskiq_deduplication.schedule import RedisDeduplicationScheduleSource class FakeScheduleSource(ScheduleSource): def __init__(self): self.startup_called = False self.shutdown_called = False self.schedules_to_return = [] self.added = [] self.deleted = [] self.pre_send_calls = [] self.post_send_calls = [] self.pre_send_raises = None async def startup(self): self.startup_called = True async def shutdown(self): self.shutdown_called = True async def get_schedules(self): return self.schedules_to_return async def add_schedule(self, schedule): self.added.append(schedule) async def delete_schedule(self, schedule_id): self.deleted.append(schedule_id) async def pre_send(self, task): self.pre_send_calls.append(task) if self.pre_send_raises is not None: raise self.pre_send_raises async def post_send(self, task): self.post_send_calls.append(task) @pytest.fixture def fake_source(): return FakeScheduleSource() @pytest.fixture def wrapper(fake_source, middleware): return RedisDeduplicationScheduleSource(fake_source, middleware) class TestDelegation: async def test_startup_delegates(self, wrapper, fake_source): await wrapper.startup() assert fake_source.startup_called async def test_shutdown_delegates(self, wrapper, fake_source): await wrapper.shutdown() assert fake_source.shutdown_called async def test_get_schedules_delegates( self, wrapper, fake_source, make_scheduled_task ): task = make_scheduled_task() fake_source.schedules_to_return = [task] assert await wrapper.get_schedules() == [task] async def test_add_schedule_delegates( self, wrapper, fake_source, make_scheduled_task ): task = make_scheduled_task() await wrapper.add_schedule(task) assert fake_source.added == [task] async def test_delete_schedule_delegates(self, wrapper, fake_source): await wrapper.delete_schedule("some-id") assert fake_source.deleted == ["some-id"] async def test_post_send_delegates(self, wrapper, fake_source, make_scheduled_task): task = make_scheduled_task() await wrapper.post_send(task) assert fake_source.post_send_calls == [task] async def test_pre_send_delegates(self, wrapper, fake_source, make_scheduled_task): task = make_scheduled_task() await wrapper.pre_send(task) assert fake_source.pre_send_calls == [task] class TestPreSend: async def test_no_lock_held_passes(self, wrapper, make_scheduled_task): task = make_scheduled_task() result = await wrapper.pre_send(task) assert result is None async def test_wrapped_source_cancellation_propagates( self, wrapper, fake_source, make_scheduled_task ): fake_source.pre_send_raises = ScheduledTaskCancelledError() with pytest.raises(ScheduledTaskCancelledError): await wrapper.pre_send(make_scheduled_task()) async def test_lock_held_raises_scheduled_task_cancelled_error( self, wrapper, middleware, make_message, make_scheduled_task ): await middleware.pre_send(make_message(task_name="my_task", kwargs={"a": 1})) task = make_scheduled_task(task_name="my_task", kwargs={"a": 1}) with pytest.raises(ScheduledTaskCancelledError): await wrapper.pre_send(task) async def test_peek_does_not_acquire_or_mutate( self, wrapper, middleware, fake_redis, make_scheduled_task ): task = make_scheduled_task(task_name="my_task", kwargs={"a": 1}) await wrapper.pre_send(task) key = middleware._build_key(task.task_name, task.labels, task.kwargs) assert not await fake_redis.exists(key) assert task.labels == {} assert task.kwargs == {"a": 1} async def test_peek_is_read_only_when_lock_held( self, wrapper, middleware, fake_redis, make_message, make_scheduled_task ): held_msg = make_message(task_name="my_task", task_id="holder", kwargs={"a": 1}) await middleware.pre_send(held_msg) key = middleware._build_deduplication_key(held_msg) ttl_before = await fake_redis.ttl(key) holder_before = await fake_redis.get(key) task = make_scheduled_task(task_name="my_task", kwargs={"a": 1}) with pytest.raises(ScheduledTaskCancelledError): await wrapper.pre_send(task) # The peek must not have re-set the key (TTL untouched) or changed # its owner. assert await fake_redis.get(key) == holder_before assert await fake_redis.ttl(key) <= ttl_before async def test_deduplication_disabled_label_bypasses_peek( self, wrapper, middleware, make_message, make_scheduled_task ): await middleware.pre_send(make_message(task_name="my_task", kwargs={"a": 1})) task = make_scheduled_task( task_name="my_task", kwargs={"a": 1}, labels={DEDUP_LABEL: False} ) await wrapper.pre_send(task) # should not raise async def test_deduplication_key_label_respected( self, wrapper, middleware, make_message, make_scheduled_task ): await middleware.pre_send( make_message(kwargs={"a": 1}, labels={DEDUP_EXPLICIT_KEY_LABEL: "fixed"}) ) task = make_scheduled_task( kwargs={"a": 999}, labels={DEDUP_EXPLICIT_KEY_LABEL: "fixed"} ) with pytest.raises(ScheduledTaskCancelledError): await wrapper.pre_send(task) async def test_deduplication_key_fields_label_respected( self, wrapper, middleware, make_message, make_scheduled_task ): await middleware.pre_send( make_message( kwargs={"a": 1, "b": 2}, labels={DEDUP_KEY_FIELDS_LABEL: ["a"]}, ) ) task = make_scheduled_task( kwargs={"a": 1, "b": 999}, labels={DEDUP_KEY_FIELDS_LABEL: ["a"]}, ) with pytest.raises(ScheduledTaskCancelledError): await wrapper.pre_send(task) async def test_non_serializable_kwargs_skips_peek_silently( self, wrapper, make_scheduled_task, caplog ): task = make_scheduled_task(kwargs={"dt": object()}) with caplog.at_level(logging.WARNING, logger="taskiq_deduplication.schedule"): await wrapper.pre_send(task) # should not raise assert not any("non-JSON-serializable" in r.message for r in caplog.records) async def test_pre_send_without_middleware_startup_raises_runtime_error( self, fake_source, make_scheduled_task ): mw = RedisDeduplicationMiddleware(redis_url="redis://localhost") w = RedisDeduplicationScheduleSource(fake_source, mw) with pytest.raises(RuntimeError, match="startup"): await w.pre_send(make_scheduled_task()) async def test_pre_send_without_middleware_startup_logs_before_raising( self, fake_source, make_scheduled_task, caplog ): mw = RedisDeduplicationMiddleware(redis_url="redis://localhost") w = RedisDeduplicationScheduleSource(fake_source, mw) with caplog.at_level(logging.ERROR, logger="taskiq_deduplication.schedule"): with pytest.raises(RuntimeError): await w.pre_send(make_scheduled_task()) assert any("startup" in r.message for r in caplog.records) async def test_different_kwargs_both_pass(self, wrapper, make_scheduled_task): await wrapper.pre_send(make_scheduled_task(kwargs={"x": 1})) await wrapper.pre_send(make_scheduled_task(kwargs={"x": 2})) class TestSchedulerIntegration: async def test_second_firing_cancelled_without_uncaught_exception( self, middleware, fake_source, make_message, make_scheduled_task ): broker = InMemoryBroker().with_middlewares(middleware) wrapper = RedisDeduplicationScheduleSource(fake_source, middleware) scheduler = TaskiqScheduler(broker=broker, sources=[wrapper]) task_name = "my_task" # Simulate a still-running first firing by holding the lock directly. await middleware.pre_send( make_message(task_name=task_name, task_id="holder", kwargs={}) ) scheduled = make_scheduled_task(task_name=task_name, kwargs={}) # Must be cancelled gracefully by on_ready()'s own except clause, not # raise DuplicateTaskError out of kiq() into an uncaught exception. await scheduler.on_ready(wrapper, scheduled)