TST: add tests for call_once

Author: plammens

pandas/tests/util/test_call_once.py

@@ -0,0 +1,84 @@
+import time
+from unittest.mock import MagicMock
+
+from pandas.util._call_once import call_once
+
+
+def test_basic():
+    callback1, callback2 = MagicMock(), MagicMock()
+    with call_once(callback1):
+        with call_once(callback1):
+            with call_once(callback2):
+                with call_once(callback1):
+                    with call_once(callback2):
+                        pass
+
+    assert callback1.call_count == 1
+    assert callback2.call_count == 1
+
+
+def test_with_key():
+    cb1, cb2 = MagicMock(), MagicMock()
+    with call_once(cb1, key="callback"):
+        with call_once(cb2, key="callback"):
+            pass
+
+    assert cb1.call_count == 1
+    assert cb2.call_count == 0
+
+
+def test_across_stack_frames():
+    callback = MagicMock()
+
+    def f():
+        with call_once(callback):
+            pass
+
+    def g():
+        with call_once(callback):
+            f()
+
+    f()
+    assert callback.call_count == 1
+    callback.reset_mock()
+
+    g()
+    assert callback.call_count == 1
+
+
+def test_concurrent_threading():
+    import threading
+
+    sleep_time = 0.01
+    callback = MagicMock()
+
+    def run(initial_sleep=0):
+        time.sleep(initial_sleep)
+        with call_once(callback):
+            with call_once(callback):
+                time.sleep(2 * sleep_time)
+
+    thread1 = threading.Thread(target=run)
+    thread2 = threading.Thread(target=run, kwargs={"initial_sleep": sleep_time})
+    thread2.start()
+    thread1.start()
+    thread1.join()
+    thread2.join()
+    assert callback.call_count == 2
+
+
+def test_concurrent_asyncio():
+    import asyncio
+
+    callback = MagicMock()
+
+    async def task():
+        with call_once(callback):
+            with call_once(callback):
+                await asyncio.sleep(0.01)
+
+    async def main():
+        await asyncio.gather(task(), task())
+
+    asyncio.run(main())
+    assert callback.call_count == 2

pandas/util/_call_once.py

@@ -21,17 +21,16 @@ class _ContextMapping(MutableMapping):
     _NO_DEFAULT_VALUE = object()
 
     def __init__(self, default_value=_NO_DEFAULT_VALUE):
-        d: dict = (
+        initial_dict: dict = (
             defaultdict(lambda: default_value)
             if default_value is not self._NO_DEFAULT_VALUE
             else dict()
         )
         # yes, we're creating a contextvar inside a closure, but it doesn't matter
         # because objects of this class will only be created at module level
         self._dict_var: contextvars.ContextVar[dict] = contextvars.ContextVar(
-            "_ContextMapping<{}>._dict_var".format(id(self))
+            "_ContextMapping<{}>._dict_var".format(id(self)), default=initial_dict
         )
-        self._dict_var.set(d)
 
     def __setitem__(self, k, v) -> None:
         d = self._dict_var.get().copy()