fix(scheduler): discover schedules of shared tasks in LabelScheduleSource

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():

tests/scheduler/test_label_based_sched.py

@@ -1,11 +1,13 @@
 import asyncio
+import logging
 from datetime import datetime, timedelta, timezone
 from typing import Any
 
 import pytest
 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()