gh-119333: Add C api to have contextvar enter/exit callbacks (#119335)

Co-authored-by: Erlend E. Aasland <erlend.aasland@protonmail.com>

1 files changed, 83 insertions, 0 deletions

diff --git a/Lib/test/test_capi/test_watchers.py b/Lib/test/test_capi/test_watchers.py
index 709b5e1..f21d262 100644
--- a/Lib/test/test_capi/test_watchers.py
+++ b/Lib/test/test_capi/test_watchers.py
@@ -1,4 +1,5 @@
 import unittest
+import contextvars
 
 from contextlib import contextmanager, ExitStack
 from test.support import (
@@ -571,5 +572,87 @@ class TestFuncWatchers(unittest.TestCase):
             _testcapi.allocate_too_many_func_watchers()
 
 
+class TestContextObjectWatchers(unittest.TestCase):
+    @contextmanager
+    def context_watcher(self, which_watcher):
+        wid = _testcapi.add_context_watcher(which_watcher)
+        try:
+            yield wid
+        finally:
+            _testcapi.clear_context_watcher(wid)
+
+    def assert_event_counts(self, exp_enter_0, exp_exit_0,
+                            exp_enter_1, exp_exit_1):
+        self.assertEqual(
+            exp_enter_0, _testcapi.get_context_watcher_num_enter_events(0))
+        self.assertEqual(
+            exp_exit_0, _testcapi.get_context_watcher_num_exit_events(0))
+        self.assertEqual(
+            exp_enter_1, _testcapi.get_context_watcher_num_enter_events(1))
+        self.assertEqual(
+            exp_exit_1, _testcapi.get_context_watcher_num_exit_events(1))
+
+    def test_context_object_events_dispatched(self):
+        # verify that all counts are zero before any watchers are registered
+        self.assert_event_counts(0, 0, 0, 0)
+
+        # verify that all counts remain zero when a context object is
+        # entered and exited with no watchers registered
+        ctx = contextvars.copy_context()
+        ctx.run(self.assert_event_counts, 0, 0, 0, 0)
+        self.assert_event_counts(0, 0, 0, 0)
+
+        # verify counts are as expected when first watcher is registered
+        with self.context_watcher(0):
+            self.assert_event_counts(0, 0, 0, 0)
+            ctx.run(self.assert_event_counts, 1, 0, 0, 0)
+            self.assert_event_counts(1, 1, 0, 0)
+
+            # again with second watcher registered
+            with self.context_watcher(1):
+                self.assert_event_counts(1, 1, 0, 0)
+                ctx.run(self.assert_event_counts, 2, 1, 1, 0)
+                self.assert_event_counts(2, 2, 1, 1)
+
+        # verify counts are reset and don't change after both watchers are cleared
+        ctx.run(self.assert_event_counts, 0, 0, 0, 0)
+        self.assert_event_counts(0, 0, 0, 0)
+
+    def test_enter_error(self):
+        with self.context_watcher(2):
+            with catch_unraisable_exception() as cm:
+                ctx = contextvars.copy_context()
+                ctx.run(int, 0)
+                self.assertEqual(
+                    cm.unraisable.err_msg,
+                    "Exception ignored in "
+                    f"Py_CONTEXT_EVENT_EXIT watcher callback for {ctx!r}"
+                )
+                self.assertEqual(str(cm.unraisable.exc_value), "boom!")
+
+    def test_exit_error(self):
+        ctx = contextvars.copy_context()
+        def _in_context(stack):
+            stack.enter_context(self.context_watcher(2))
+
+        with catch_unraisable_exception() as cm:
+            with ExitStack() as stack:
+                ctx.run(_in_context, stack)
+            self.assertEqual(str(cm.unraisable.exc_value), "boom!")
+
+    def test_clear_out_of_range_watcher_id(self):
+        with self.assertRaisesRegex(ValueError, r"Invalid context watcher ID -1"):
+            _testcapi.clear_context_watcher(-1)
+        with self.assertRaisesRegex(ValueError, r"Invalid context watcher ID 8"):
+            _testcapi.clear_context_watcher(8)  # CONTEXT_MAX_WATCHERS = 8
+
+    def test_clear_unassigned_watcher_id(self):
+        with self.assertRaisesRegex(ValueError, r"No context watcher set for ID 1"):
+            _testcapi.clear_context_watcher(1)
+
+    def test_allocate_too_many_watchers(self):
+        with self.assertRaisesRegex(RuntimeError, r"no more context watcher IDs available"):
+            _testcapi.allocate_too_many_context_watchers()
+
 if __name__ == "__main__":
     unittest.main()