diff --git a/taskiq/schedule_sources/label_based.py b/taskiq/schedule_sources/label_based.py index 663e1db0..fa53624f 100644 --- a/taskiq/schedule_sources/label_based.py +++ b/taskiq/schedule_sources/label_based.py @@ -34,14 +34,18 @@ async def startup(self) -> None: """ self.schedules.clear() for task_name, task in self.broker.get_all_tasks().items(): - if task.broker != self.broker: - # if task broker doesn't match self, something is probably wrong - logger.warning( - "Broker for %s `%s` doesn't match scheduler's broker `%s`", - task_name, - task.broker, - self.broker, - ) + # Shared tasks are bound to the shared broker, + # but they are sent using its default broker. + task_broker = getattr(task.broker, "_default_broker", None) or task.broker + if task_broker != self.broker: + if task.labels.get("schedule"): + # task has schedules, but they are going to be ignored + logger.warning( + "Broker for %s `%s` doesn't match scheduler's broker `%s`", + task_name, + task_broker, + self.broker, + ) continue for schedule in task.labels.get("schedule", []): if not {"cron", "interval", "time"} & schedule.keys(): diff --git a/tests/scheduler/test_label_based_sched.py b/tests/scheduler/test_label_based_sched.py index 81f5edfe..d466e648 100644 --- a/tests/scheduler/test_label_based_sched.py +++ b/tests/scheduler/test_label_based_sched.py @@ -1,4 +1,5 @@ import asyncio +import logging from datetime import datetime, timedelta, timezone from typing import Any @@ -6,6 +7,7 @@ from freezegun import freeze_time from taskiq.brokers.inmemory_broker import InMemoryBroker +from taskiq.brokers.shared_broker import AsyncSharedBroker from taskiq.cli.scheduler.args import SchedulerArgs from taskiq.cli.scheduler.run import run_scheduler from taskiq.schedule_sources.label_based import LabelScheduleSource @@ -74,6 +76,69 @@ def task() -> None: assert schedules == [] +async def test_label_discovery_shared_task( + caplog: pytest.LogCaptureFixture, +) -> None: + broker = InMemoryBroker() + shared_broker = AsyncSharedBroker() + shared_broker.default_broker(broker) + + @shared_broker.task(task_name="shared_task", schedule=[{"cron": "* * * * *"}]) + def shared_task() -> None: + pass + + @shared_broker.task(task_name="shared_task_no_schedule") + def shared_task_no_schedule() -> None: + pass + + source = LabelScheduleSource(broker) + with caplog.at_level(logging.WARNING): + await source.startup() + + schedules = await source.get_schedules() + assert [schedule.task_name for schedule in schedules] == ["shared_task"] + assert "doesn't match scheduler's broker" not in caplog.text + + +@pytest.mark.parametrize( + "set_default_broker", + [ + pytest.param(True, id="other_default_broker"), + pytest.param(False, id="no_default_broker"), + ], +) +async def test_label_discovery_shared_task_broker_mismatch( + caplog: pytest.LogCaptureFixture, + set_default_broker: bool, +) -> None: + broker = InMemoryBroker() + shared_broker = AsyncSharedBroker() + if set_default_broker: + shared_broker.default_broker(InMemoryBroker()) + + @shared_broker.task(task_name="shared_task", schedule=[{"cron": "* * * * *"}]) + def shared_task() -> None: + pass + + @shared_broker.task(task_name="shared_task_no_schedule") + def shared_task_no_schedule() -> None: + pass + + source = LabelScheduleSource(broker) + with caplog.at_level(logging.WARNING): + await source.startup() + + assert await source.get_schedules() == [] + # Only the task whose schedules are ignored is reported. + warnings = [ + record.getMessage() + for record in caplog.records + if "doesn't match scheduler's broker" in record.getMessage() + ] + assert len(warnings) == 1 + assert warnings[0].startswith("Broker for shared_task `") + + async def test_task_scheduled_at_time_runs_only_once(mock_sleep: None) -> None: event = asyncio.Event() broker = InMemoryBroker()