test decorator

Author: sanderegg

packages/service-library/src/servicelib/background_task.py

@@ -4,15 +4,9 @@
 import functools
 import logging
 from collections.abc import AsyncIterator, Awaitable, Callable
-from typing import Final, ParamSpec, TypeVar
-
-from tenacity import (
-    TryAgain,
-    before_sleep_log,
-    retry,
-    retry_always,
-    retry_if_exception_type,
-)
+from typing import Any, Coroutine, Final, ParamSpec, TypeVar
+
+from tenacity import TryAgain, before_sleep_log, retry, retry_if_exception_type
 from tenacity.asyncio import AsyncRetrying
 from tenacity.wait import wait_fixed
 
@@ -46,8 +40,12 @@ def periodic(
     interval: datetime.timedelta,
     raise_on_error: bool = False,
     early_wake_up_event: asyncio.Event | None = None,
-) -> Callable[[Callable[P, Awaitable[None]]], Callable[P, Awaitable[None]]]:
-    def _decorator(func: Callable[P, Awaitable[None]]) -> Callable[P, Awaitable[None]]:
+) -> Callable[
+    [Callable[P, Coroutine[Any, Any, None]]], Callable[P, Coroutine[Any, Any, None]]
+]:
+    def _decorator(
+        func: Callable[P, Coroutine[Any, Any, None]],
+    ) -> Callable[P, Coroutine[Any, Any, None]]:
         nap = (
             asyncio.sleep
             if early_wake_up_event is None
@@ -99,7 +97,9 @@ async def _periodic_scheduled_task(
         sleep=nap,
         wait=wait_fixed(interval.total_seconds()),
         reraise=True,
-        retry=retry_if_exception_type(TryAgain) if raise_on_error else retry_always,
+        retry=retry_if_exception_type(TryAgain)
+        if raise_on_error
+        else retry_if_exception_type(),
     ):
         with attempt:
             with (

packages/service-library/tests/test_background_task.py

@@ -9,11 +9,13 @@
 from collections.abc import AsyncIterator, Awaitable, Callable
 from typing import Final
 from unittest import mock
+from unittest.mock import AsyncMock
 
 import pytest
 from faker import Faker
 from pytest_mock.plugin import MockerFixture
-from servicelib.background_task import (
+from servicelib.background_task import (  # Assuming the module is imported correctly
+    periodic,
     periodic_task,
     start_periodic_task,
     stop_periodic_task,
@@ -177,3 +179,24 @@ async def test_periodic_task_context_manager(
         assert asyncio_task.cancelled() is False
         assert asyncio_task.done() is False
     assert asyncio_task.cancelled() is True
+
+
+async def test_periodic_decorator():
+    # This mock function will allow us to test if the function is called periodically
+    mock_func = AsyncMock()
+
+    @periodic(interval=datetime.timedelta(seconds=0.1))
+    async def _func() -> None:
+        await mock_func()
+
+    task = asyncio.create_task(_func())
+
+    # Give some time for the periodic calls to happen
+    await asyncio.sleep(0.5)
+
+    # Once enough time has passed, cancel the task
+    task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await task
+
+    assert mock_func.call_count > 1