diff --git a/tests/fixtures/time.py b/tests/fixtures/time.py index 403c3dccb661..538a7594b590 100644 --- a/tests/fixtures/time.py +++ b/tests/fixtures/time.py @@ -1,7 +1,9 @@ -from typing import Optional, Union +from datetime import timedelta +from typing import Callable, Optional, Union import pendulum import pytest +from pendulum import DateTime from pendulum.tz.timezone import Timezone @@ -16,3 +18,27 @@ def frozen_time(tz: Optional[Union[str, Timezone]] = None): monkeypatch.setattr(pendulum, "now", frozen_time) return frozen + + +@pytest.fixture +def advance_time(monkeypatch: pytest.MonkeyPatch) -> Callable[[timedelta], DateTime]: + clock = pendulum.now("UTC") + + def advance(amount: timedelta): + nonlocal clock + clock += amount + return clock + + def nowish(tz: Optional[Union[str, Timezone]] = None): + # each time this is called, advance by 1 microsecond so that time is moving + # forward bit-by-bit to avoid everything appearing to happen all at once + advance(timedelta(microseconds=1)) + + if tz is None: + return clock + + return clock.in_timezone(tz) + + monkeypatch.setattr(pendulum, "now", nowish) + + return advance diff --git a/tests/test_task_engine.py b/tests/test_task_engine.py index 4aa7d839bdec..64c2e235ca0b 100644 --- a/tests/test_task_engine.py +++ b/tests/test_task_engine.py @@ -1139,7 +1139,9 @@ async def async_task(): assert await state.result() == 1800 assert state.data.storage_key == "foo-bar" - async def test_cache_expiration_is_respected(self, prefect_client, tmp_path): + async def test_cache_expiration_is_respected( + self, prefect_client, tmp_path, advance_time + ): fs = LocalFileSystem(basepath=tmp_path) @task( @@ -1164,7 +1166,7 @@ async def async_task(): assert first_result == second_result, "Cache was not used" # let cache expire... - await asyncio.sleep(1.1) + advance_time(timedelta(seconds=1.1)) third_state = await async_task(return_state=True) assert third_state.is_completed()