Skip to content

Commit

Permalink
Only mock asyncio on request
Browse files Browse the repository at this point in the history
  • Loading branch information
bblommers committed Dec 7, 2023
1 parent 9152728 commit 7ca707d
Show file tree
Hide file tree
Showing 4 changed files with 28 additions and 21 deletions.
41 changes: 22 additions & 19 deletions freezegun/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -544,7 +544,7 @@ def move_to(self, target_datetime):

class _freeze_time:

def __init__(self, time_to_freeze_str, tz_offset, ignore, tick, as_arg, as_kwarg, auto_tick_seconds):
def __init__(self, time_to_freeze_str, tz_offset, ignore, tick, as_arg, as_kwarg, auto_tick_seconds, real_asyncio):
self.time_to_freeze = _parse_time_to_freeze(time_to_freeze_str)
self.tz_offset = _parse_tz_offset(tz_offset)
self.ignore = tuple(ignore)
Expand All @@ -554,6 +554,7 @@ def __init__(self, time_to_freeze_str, tz_offset, ignore, tick, as_arg, as_kwarg
self.modules_at_start = set()
self.as_arg = as_arg
self.as_kwarg = as_kwarg
self.real_asyncio = real_asyncio

def __call__(self, func):
if inspect.isclass(func):
Expand Down Expand Up @@ -727,20 +728,21 @@ def start(self):
setattr(module, attribute_name, fake)
add_change((module, attribute_name, attribute_value))

# To avoid breaking `asyncio.sleep()`, let asyncio event loops see real
# monotonic time even though we've just frozen `time.monotonic()` which
# is normally used there. If we didn't do this, `await asyncio.sleep()`
# would be hanging forever breaking many tests that use `freeze_time`.
#
# Note that we cannot statically tell the class of asyncio event loops
# because it is not officially documented and can actually be changed
# at run time using `asyncio.set_event_loop_policy`. That's why we check
# the type by creating a loop here and destroying it immediately.
event_loop = asyncio.new_event_loop()
event_loop.close()
EventLoopClass = type(event_loop)
add_change((EventLoopClass, "time", EventLoopClass.time))
EventLoopClass.time = lambda self: real_monotonic()
if self.real_asyncio:
# To avoid breaking `asyncio.sleep()`, let asyncio event loops see real
# monotonic time even though we've just frozen `time.monotonic()` which
# is normally used there. If we didn't do this, `await asyncio.sleep()`
# would be hanging forever breaking many tests that use `freeze_time`.
#
# Note that we cannot statically tell the class of asyncio event loops
# because it is not officially documented and can actually be changed
# at run time using `asyncio.set_event_loop_policy`. That's why we check
# the type by creating a loop here and destroying it immediately.
event_loop = asyncio.new_event_loop()
event_loop.close()
EventLoopClass = type(event_loop)
add_change((EventLoopClass, "time", EventLoopClass.time))
EventLoopClass.time = lambda self: real_monotonic()

return freeze_factory

Expand Down Expand Up @@ -830,7 +832,7 @@ def wrapper(*args, **kwargs):


def freeze_time(time_to_freeze=None, tz_offset=0, ignore=None, tick=False, as_arg=False, as_kwarg='',
auto_tick_seconds=0):
auto_tick_seconds=0, real_asyncio=False):
acceptable_times = (type(None), str, datetime.date, datetime.timedelta,
types.FunctionType, types.GeneratorType)

Expand All @@ -845,14 +847,14 @@ def freeze_time(time_to_freeze=None, tz_offset=0, ignore=None, tick=False, as_ar
raise SystemError('Calling freeze_time with tick=True is only compatible with CPython')

if isinstance(time_to_freeze, types.FunctionType):
return freeze_time(time_to_freeze(), tz_offset, ignore, tick, as_arg, as_kwarg, auto_tick_seconds)
return freeze_time(time_to_freeze(), tz_offset, ignore, tick, as_arg, as_kwarg, auto_tick_seconds, real_asyncio=real_asyncio)

if isinstance(time_to_freeze, types.GeneratorType):
return freeze_time(next(time_to_freeze), tz_offset, ignore, tick, as_arg, as_kwarg, auto_tick_seconds)
return freeze_time(next(time_to_freeze), tz_offset, ignore, tick, as_arg, as_kwarg, auto_tick_seconds, real_asyncio=real_asyncio)

if MayaDT is not None and isinstance(time_to_freeze, MayaDT):
return freeze_time(time_to_freeze.datetime(), tz_offset, ignore,
tick, as_arg, as_kwarg, auto_tick_seconds)
tick, as_arg, as_kwarg, auto_tick_seconds, real_asyncio=real_asyncio)

if ignore is None:
ignore = []
Expand All @@ -868,6 +870,7 @@ def freeze_time(time_to_freeze=None, tz_offset=0, ignore=None, tick=False, as_ar
as_arg=as_arg,
as_kwarg=as_kwarg,
auto_tick_seconds=auto_tick_seconds,
real_asyncio=real_asyncio,
)


Expand Down
2 changes: 2 additions & 0 deletions freezegun/api.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ class _freeze_time:
as_arg: bool,
as_kwarg: str,
auto_tick_seconds: float,
real_asyncio: bool,
) -> None: ...
@overload
def __call__(self, func: type[_T]) -> type[_T]: ...
Expand All @@ -57,4 +58,5 @@ def freeze_time(
as_arg: bool | None = ...,
as_kwarg: str | None = ...,
auto_tick_seconds: float | None = ...,
real_asyncio: bool | None = ...
) -> _freeze_time: ...
4 changes: 2 additions & 2 deletions tests/test_asyncio.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@ def test_asyncio_sleeping_not_affected_by_freeze_time():
async def coroutine():
# Sleeping with time frozen should sleep the expected duration.
before_sleep = time.time()
with freeze_time('1970-01-02'):
with freeze_time('1970-01-02', real_asyncio=True):
await asyncio.sleep(0.05)
assert 0.02 <= time.time() - before_sleep < 0.3

Expand Down Expand Up @@ -76,5 +76,5 @@ async def coroutine():
await asyncio.sleep(0.15)
assert timestamps == [86400]

with freeze_time('1970-01-02'):
with freeze_time('1970-01-02', real_asyncio=True):
asyncio.run(coroutine())
2 changes: 2 additions & 0 deletions tests/test_configure.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ def test_default_ignore_list_is_overridden():
as_arg=False,
as_kwarg='',
auto_tick_seconds=0,
real_asyncio=False,
)

def test_extend_default_ignore_list():
Expand Down Expand Up @@ -64,4 +65,5 @@ def test_extend_default_ignore_list():
as_arg=False,
as_kwarg='',
auto_tick_seconds=0,
real_asyncio=False,
)

0 comments on commit 7ca707d

Please sign in to comment.