Merge pull request #470 from marcinsulikowski/improve-asyncio-support

Improve asyncio support to avoid hangs of asyncio.sleep()

AUTHORS.rst

@@ -20,3 +20,4 @@ Patches and Suggestions
 - `Lukasz Balcerzak <https://github.com/lukaszb>`_
 - `Hannes Ljungberg <hannes@5monkeys.se>`_
 - `staticdev <staticdev-support@proton.me>`_
+- `Marcin Sulikowski <https://github.com/marcinsulikowski>`_

CHANGELOG

@@ -1,6 +1,11 @@
 Freezegun Changelog
 ===================
 
+1.3.0
+-----
+
+* Fixed `asyncio` support to avoid `await asyncio.sleep(1)` hanging forever.
+
 1.2.2
 -----
 

freezegun/api.py

@@ -1,5 +1,6 @@
 from . import config
 from ._async import wrap_coroutine
+import asyncio
 import copyreg
 import dateutil
 import datetime
@@ -726,6 +727,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()
+
         return freeze_factory
 
     def stop(self):
@@ -739,8 +755,8 @@ def stop(self):
             datetime.date = real_date
             copyreg.dispatch_table.pop(real_datetime)
             copyreg.dispatch_table.pop(real_date)
-            for module, module_attribute, original_value in self.undo_changes:
-                setattr(module, module_attribute, original_value)
+            for module_or_object, attribute, original_value in self.undo_changes:
+                setattr(module_or_object, attribute, original_value)
             self.undo_changes = []
 
             # Restore modules loaded after start()

tests/test_asyncio.py

@@ -1,17 +1,80 @@
 import asyncio
 import datetime
-from textwrap import dedent
-from unittest import SkipTest
+import time
 
 from freezegun import freeze_time
 
 
-def test_time_freeze_coroutine():
-    if not asyncio:
-        raise SkipTest('asyncio required')
-
+def test_datetime_in_coroutine():
     @freeze_time('1970-01-01')
     async def frozen_coroutine():
         assert datetime.date.today() == datetime.date(1970, 1, 1)
 
     asyncio.run(frozen_coroutine())
+
+
+def test_freezing_time_in_coroutine():
+    """Test calling freeze_time while executing asyncio loop."""
+    async def coroutine():
+        with freeze_time('1970-01-02'):
+            assert time.time() == 86400
+        with freeze_time('1970-01-03'):
+            assert time.time() == 86400 * 2
+
+    asyncio.run(coroutine())
+
+
+def test_freezing_time_before_running_coroutine():
+    """Test calling freeze_time before executing asyncio loop."""
+    async def coroutine():
+        assert time.time() == 86400
+    with freeze_time('1970-01-02'):
+        asyncio.run(coroutine())
+
+
+def test_asyncio_sleeping_not_affected_by_freeze_time():
+    """Test that asyncio.sleep() is not affected by `freeze_time`.
+
+    This test ensures that despite freezing time using `freeze_time`,
+    the asyncio event loop can see real monotonic time, which is required
+    to make things like `asyncio.sleep()` work.
+    """
+
+    async def coroutine():
+        # Sleeping with time frozen should sleep the expected duration.
+        before_sleep = time.time()
+        with freeze_time('1970-01-02'):
+            await asyncio.sleep(0.05)
+        assert 0.02 <= time.time() - before_sleep < 0.3
+
+        # Exiting `freeze_time` the time should not break asyncio sleeping.
+        before_sleep = time.time()
+        await asyncio.sleep(0.05)
+        assert 0.02 <= time.time() - before_sleep < 0.3
+
+    asyncio.run(coroutine())
+
+
+def test_asyncio_to_call_later_with_frozen_time():
+    """Test that asyncio `loop.call_later` works with frozen time."""
+    # `to_call_later` will be called by asyncio event loop and should add
+    # the Unix timestamp of 1970-01-02 00:00 to the `timestamps` list.
+    timestamps = []
+    def to_call_later():
+        timestamps.append(time.time())
+
+    async def coroutine():
+        # Schedule calling `to_call_later` in 100 ms.
+        asyncio.get_running_loop().call_later(0.1, to_call_later)
+
+        # Sleeping for 10 ms should not result in calling `to_call_later`.
+        await asyncio.sleep(0.01)
+        assert timestamps == []
+
+        # But sleeping more (150 ms in this case) should call `to_call_later`
+        # and we should see `timestamps` updated.
+        await asyncio.sleep(0.15)
+        assert timestamps == [86400]
+
+    with freeze_time('1970-01-02'):
+        asyncio.run(coroutine())