Feature/label tests (#56)

* Added test for retry middleware.
* Added test to schedules.

Signed-off-by: Pavel Kirilin <[email protected]>

Author: s3rius

taskiq/kicker.py

@@ -92,15 +92,15 @@ async def kiq(  # noqa: D102
         self: "AsyncKicker[_FuncParams, Coroutine[Any, Any, _T]]",
         *args: _FuncParams.args,
         **kwargs: _FuncParams.kwargs,
-    ) -> AsyncTaskiqTask[_T]:
+    ) -> AsyncTaskiqTask[_T]:  # pragma: no cover
         ...
 
     @overload
     async def kiq(  # noqa: D102
         self: "AsyncKicker[_FuncParams, _ReturnType]",
         *args: _FuncParams.args,
         **kwargs: _FuncParams.kwargs,
-    ) -> AsyncTaskiqTask[_ReturnType]:
+    ) -> AsyncTaskiqTask[_ReturnType]:  # pragma: no cover
         ...
 
     async def kiq(  # noqa: C901

taskiq/scheduler/merge_functions.py

@@ -34,9 +34,8 @@ def only_unique(
     :param new_tasks: newly discovered tasks.
     :return: list of unique schedules.
     """
-    new_tasks = []
-    new_tasks.extend(old_tasks)
+    result = old_tasks
     for task in new_tasks:
-        if task not in new_tasks:
-            new_tasks.append(task)
-    return new_tasks
+        if task not in result:
+            result.append(task)
+    return result

taskiq/scheduler/scheduler.py

@@ -31,13 +31,13 @@ def __init__(
             List["ScheduledTask"],
         ] = preserve_all,
         refresh_delay: float = 30.0,
-    ) -> None:
+    ) -> None:  # pragma: no cover
         self.broker = broker
         self.sources = sources
         self.refresh_delay = refresh_delay
         self.merge_func = merge_func
 
-    async def startup(self) -> None:
+    async def startup(self) -> None:  # pragma: no cover
         """
         This method is called on startup.
 

tests/scheduler/test_label_based_sched.py

@@ -0,0 +1,45 @@
+import pytest
+
+from taskiq.brokers.inmemory_broker import InMemoryBroker
+from taskiq.schedule_sources.label_based import LabelScheduleSource
+from taskiq.scheduler.scheduler import ScheduledTask
+
+
+@pytest.mark.anyio
+async def test_label_discovery() -> None:
+    broker = InMemoryBroker()
+
+    @broker.task(
+        task_name="test_task",
+        schedule=[{"cron": "* * * * *"}],
+    )
+    def task() -> None:
+        pass
+
+    source = LabelScheduleSource(broker)
+    schedules = await source.get_schedules()
+    assert schedules == [
+        ScheduledTask(
+            cron="* * * * *",
+            task_name="test_task",
+            labels={"schedule": [{"cron": "* * * * *"}]},
+            args=[],
+            kwargs={},
+        ),
+    ]
+
+
+@pytest.mark.anyio
+async def test_label_discovery_no_cron() -> None:
+    broker = InMemoryBroker()
+
+    @broker.task(
+        task_name="test_task",
+        schedule=[{"args": ["* * * * *"]}],
+    )
+    def task() -> None:
+        pass
+
+    source = LabelScheduleSource(broker)
+    schedules = await source.get_schedules()
+    assert schedules == []

tests/scheduler/test_merge_funcs.py

@@ -0,0 +1,13 @@
+from taskiq.scheduler.merge_functions import only_unique, preserve_all
+
+
+def test_preserve_all() -> None:
+    first = [1, 2, 3]
+    second = [3, 4, 5]
+    assert preserve_all(first, second) == [1, 2, 3, 3, 4, 5]  # type: ignore
+
+
+def test_only_unique() -> None:
+    first = [1, 2, 3]
+    second = [3, 4, 5]
+    assert only_unique(first, second) == [1, 2, 3, 4, 5]  # type: ignore