change schedules storage to list of tuples

Author: so-saf

taskiq/cli/scheduler/run.py

@@ -3,7 +3,7 @@
 import sys
 from datetime import datetime, timedelta, timezone
 from logging import basicConfig, getLogger
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 from zoneinfo import ZoneInfo
 
 import pycron
@@ -59,23 +59,22 @@ async def get_schedules(source: ScheduleSource) -> List[ScheduledTask]:
 
 async def get_all_schedules(
     scheduler: TaskiqScheduler,
-) -> Dict[ScheduleSource, List[ScheduledTask]]:
+) -> List[Tuple[ScheduleSource, List[ScheduledTask]]]:
     """
     Task to update all schedules.
 
     This function updates all schedules
-    from all sources and returns a dict
-    with source as a key and list of
-    scheduled tasks as a value.
+    from all sources and returns a list
+    of (source, tasks) pairs.
 
     :param scheduler: current scheduler.
-    :return: dict with source as a key and list of scheduled tasks as a value.
+    :return: list of (source, tasks) pairs.
     """
     logger.debug("Started schedule update.")
     schedules: List[List[ScheduledTask]] = await asyncio.gather(
         *[get_schedules(source) for source in scheduler.sources],
     )
-    return dict(zip(scheduler.sources, schedules))
+    return list(zip(scheduler.sources, schedules))
 
 
 class CronValueError(Exception):
@@ -197,15 +196,15 @@ def __init__(
         self.interval_tasks_last_run: dict[ScheduleId, datetime] = {}
         self.time_tasks_last_run: dict[ScheduleId, datetime] = {}
 
-        self.scheduled_tasks: Dict[ScheduleSource, List[ScheduledTask]] = {}
+        self.scheduled_tasks: List[Tuple[ScheduleSource, List[ScheduledTask]]] = []
         self.scheduled_tasks_updated_at: Optional[datetime] = None
         self._update_schedules_task_future: Optional[asyncio.Task[Any]] = None
 
     def _update_schedules_task_future_callback(self, task_: asyncio.Task[Any]) -> None:
         self.scheduled_tasks = task_.result()
 
         new_schedules_ids: set[ScheduleId] = set()
-        for source, task_list in self.scheduled_tasks.items():
+        for source, task_list in self.scheduled_tasks:
             logger.debug("Got %d schedules from source %s.", len(task_list), source)
             new_schedules_ids.update({t.schedule_id for t in task_list})
 
@@ -237,7 +236,7 @@ async def _update_scheduled_tasks(self) -> None:
 
     def _mark_cron_tasks_as_already_run(self) -> None:
         current_minute = datetime.now(tz=timezone.utc).replace(second=0, microsecond=0)
-        for _, task_list in self.scheduled_tasks.items():
+        for _, task_list in self.scheduled_tasks:
             for task in task_list:
                 if task.cron is not None:
                     self.cron_tasks_last_run[task.schedule_id] = current_minute
@@ -328,7 +327,7 @@ async def run(
                 await self._update_scheduled_tasks()
                 self.scheduled_tasks_updated_at = now
 
-            for source, task_list in self.scheduled_tasks.items():
+            for source, task_list in self.scheduled_tasks:
                 for task in task_list:
                     is_ready_to_send: bool = self._is_schedule_ready_to_send(
                         task=task,

tests/cli/scheduler/test_updater.py

@@ -53,10 +53,10 @@ async def test_get_schedules_success() -> None:
     schedules = await get_all_schedules(
         TaskiqScheduler(InMemoryBroker(), sources),
     )
-    assert schedules == {
-        sources[0]: schedules1,
-        sources[1]: schedules2,
-    }
+    assert schedules == [
+        (sources[0], schedules1),
+        (sources[1], schedules2),
+    ]
 
 
 async def test_get_schedules_error() -> None:
@@ -77,7 +77,7 @@ async def test_get_schedules_error() -> None:
     schedules = await get_all_schedules(
         TaskiqScheduler(InMemoryBroker(), [source1, source2]),
     )
-    assert schedules == {
-        source1: source1.schedules,
-        source2: [],
-    }
+    assert schedules == [
+        (source1, source1.schedules),
+        (source2, []),
+    ]