test_contextvars.py 9.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266
  1. import unittest
  2. import gc
  3. import sys
  4. from functools import partial
  5. from greenlet import greenlet
  6. from greenlet import getcurrent
  7. try:
  8. from contextvars import Context
  9. from contextvars import ContextVar
  10. from contextvars import copy_context
  11. except ImportError:
  12. Context = ContextVar = copy_context = None
  13. # We don't support testing if greenlet's built-in context var support is disabled.
  14. @unittest.skipUnless(Context is not None, "ContextVar not supported")
  15. class ContextVarsTests(unittest.TestCase):
  16. def _new_ctx_run(self, *args, **kwargs):
  17. return copy_context().run(*args, **kwargs)
  18. def _increment(self, greenlet_id, ctx_var, callback, counts, expect):
  19. if expect is None:
  20. self.assertIsNone(ctx_var.get())
  21. else:
  22. self.assertEqual(ctx_var.get(), expect)
  23. ctx_var.set(greenlet_id)
  24. for _ in range(2):
  25. counts[ctx_var.get()] += 1
  26. callback()
  27. def _test_context(self, propagate_by):
  28. id_var = ContextVar("id", default=None)
  29. id_var.set(0)
  30. callback = getcurrent().switch
  31. counts = dict((i, 0) for i in range(5))
  32. lets = [
  33. greenlet(partial(
  34. partial(
  35. copy_context().run,
  36. self._increment
  37. ) if propagate_by == "run" else self._increment,
  38. greenlet_id=i,
  39. ctx_var=id_var,
  40. callback=callback,
  41. counts=counts,
  42. expect=(
  43. i - 1 if propagate_by == "share" else
  44. 0 if propagate_by in ("set", "run") else None
  45. )
  46. ))
  47. for i in range(1, 5)
  48. ]
  49. for let in lets:
  50. if propagate_by == "set":
  51. let.gr_context = copy_context()
  52. elif propagate_by == "share":
  53. let.gr_context = getcurrent().gr_context
  54. for i in range(2):
  55. counts[id_var.get()] += 1
  56. for let in lets:
  57. let.switch()
  58. if propagate_by == "run":
  59. # Must leave each context.run() in reverse order of entry
  60. for let in reversed(lets):
  61. let.switch()
  62. else:
  63. # No context.run(), so fine to exit in any order.
  64. for let in lets:
  65. let.switch()
  66. for let in lets:
  67. self.assertTrue(let.dead)
  68. # When using run(), we leave the run() as the greenlet dies,
  69. # and there's no context "underneath". When not using run(),
  70. # gr_context still reflects the context the greenlet was
  71. # running in.
  72. self.assertEqual(let.gr_context is None, propagate_by == "run")
  73. if propagate_by == "share":
  74. self.assertEqual(counts, {0: 1, 1: 1, 2: 1, 3: 1, 4: 6})
  75. else:
  76. self.assertEqual(set(counts.values()), set([2]))
  77. def test_context_propagated_by_context_run(self):
  78. self._new_ctx_run(self._test_context, "run")
  79. def test_context_propagated_by_setting_attribute(self):
  80. self._new_ctx_run(self._test_context, "set")
  81. def test_context_not_propagated(self):
  82. self._new_ctx_run(self._test_context, None)
  83. def test_context_shared(self):
  84. self._new_ctx_run(self._test_context, "share")
  85. def test_break_ctxvars(self):
  86. let1 = greenlet(copy_context().run)
  87. let2 = greenlet(copy_context().run)
  88. let1.switch(getcurrent().switch)
  89. let2.switch(getcurrent().switch)
  90. # Since let2 entered the current context and let1 exits its own, the
  91. # interpreter emits:
  92. # RuntimeError: cannot exit context: thread state references a different context object
  93. let1.switch()
  94. def test_not_broken_if_using_attribute_instead_of_context_run(self):
  95. let1 = greenlet(getcurrent().switch)
  96. let2 = greenlet(getcurrent().switch)
  97. let1.gr_context = copy_context()
  98. let2.gr_context = copy_context()
  99. let1.switch()
  100. let2.switch()
  101. let1.switch()
  102. let2.switch()
  103. def test_context_assignment_while_running(self):
  104. id_var = ContextVar("id", default=None)
  105. def target():
  106. self.assertIsNone(id_var.get())
  107. self.assertIsNone(gr.gr_context)
  108. # Context is created on first use
  109. id_var.set(1)
  110. self.assertIsInstance(gr.gr_context, Context)
  111. self.assertEqual(id_var.get(), 1)
  112. self.assertEqual(gr.gr_context[id_var], 1)
  113. # Clearing the context makes it get re-created as another
  114. # empty context when next used
  115. old_context = gr.gr_context
  116. gr.gr_context = None # assign None while running
  117. self.assertIsNone(id_var.get())
  118. self.assertIsNone(gr.gr_context)
  119. id_var.set(2)
  120. self.assertIsInstance(gr.gr_context, Context)
  121. self.assertEqual(id_var.get(), 2)
  122. self.assertEqual(gr.gr_context[id_var], 2)
  123. new_context = gr.gr_context
  124. getcurrent().parent.switch((old_context, new_context))
  125. # parent switches us back to old_context
  126. self.assertEqual(id_var.get(), 1)
  127. gr.gr_context = new_context # assign non-None while running
  128. self.assertEqual(id_var.get(), 2)
  129. getcurrent().parent.switch()
  130. # parent switches us back to no context
  131. self.assertIsNone(id_var.get())
  132. self.assertIsNone(gr.gr_context)
  133. gr.gr_context = old_context
  134. self.assertEqual(id_var.get(), 1)
  135. getcurrent().parent.switch()
  136. # parent switches us back to no context
  137. self.assertIsNone(id_var.get())
  138. self.assertIsNone(gr.gr_context)
  139. gr = greenlet(target)
  140. with self.assertRaisesRegex(AttributeError, "can't delete attr"):
  141. del gr.gr_context
  142. self.assertIsNone(gr.gr_context)
  143. old_context, new_context = gr.switch()
  144. self.assertIs(new_context, gr.gr_context)
  145. self.assertEqual(old_context[id_var], 1)
  146. self.assertEqual(new_context[id_var], 2)
  147. self.assertEqual(new_context.run(id_var.get), 2)
  148. gr.gr_context = old_context # assign non-None while suspended
  149. gr.switch()
  150. self.assertIs(gr.gr_context, new_context)
  151. gr.gr_context = None # assign None while suspended
  152. gr.switch()
  153. self.assertIs(gr.gr_context, old_context)
  154. gr.gr_context = None
  155. gr.switch()
  156. self.assertIsNone(gr.gr_context)
  157. # Make sure there are no reference leaks
  158. gr = None
  159. gc.collect()
  160. self.assertEqual(sys.getrefcount(old_context), 2)
  161. self.assertEqual(sys.getrefcount(new_context), 2)
  162. def test_context_assignment_different_thread(self):
  163. import threading
  164. ctx = Context()
  165. var = ContextVar("var", default=None)
  166. is_running = threading.Event()
  167. should_suspend = threading.Event()
  168. did_suspend = threading.Event()
  169. should_exit = threading.Event()
  170. holder = []
  171. def greenlet_in_thread_fn():
  172. var.set(1)
  173. is_running.set()
  174. should_suspend.wait()
  175. var.set(2)
  176. getcurrent().parent.switch()
  177. holder.append(var.get())
  178. def thread_fn():
  179. gr = greenlet(greenlet_in_thread_fn)
  180. gr.gr_context = ctx
  181. holder.append(gr)
  182. gr.switch()
  183. did_suspend.set()
  184. should_exit.wait()
  185. gr.switch()
  186. thread = threading.Thread(target=thread_fn, daemon=True)
  187. thread.start()
  188. is_running.wait()
  189. gr = holder[0]
  190. # Can't access or modify context if the greenlet is running
  191. # in a different thread
  192. with self.assertRaisesRegex(ValueError, "running in a different"):
  193. getattr(gr, 'gr_context')
  194. with self.assertRaisesRegex(ValueError, "running in a different"):
  195. gr.gr_context = None
  196. should_suspend.set()
  197. did_suspend.wait()
  198. # OK to access and modify context if greenlet is suspended
  199. self.assertIs(gr.gr_context, ctx)
  200. self.assertEqual(gr.gr_context[var], 2)
  201. gr.gr_context = None
  202. should_exit.set()
  203. thread.join()
  204. self.assertEqual(holder, [gr, None])
  205. # Context can still be accessed/modified when greenlet is dead:
  206. self.assertIsNone(gr.gr_context)
  207. gr.gr_context = ctx
  208. self.assertIs(gr.gr_context, ctx)
  209. @unittest.skipIf(Context is not None, "ContextVar supported")
  210. class NoContextVarsTests(unittest.TestCase):
  211. def test_contextvars_errors(self):
  212. let1 = greenlet(getcurrent().switch)
  213. self.assertFalse(hasattr(let1, 'gr_context'))
  214. with self.assertRaises(AttributeError):
  215. getattr(let1, 'gr_context')
  216. with self.assertRaises(AttributeError):
  217. let1.gr_context = None
  218. let1.switch()
  219. with self.assertRaises(AttributeError):
  220. getattr(let1, 'gr_context')
  221. with self.assertRaises(AttributeError):
  222. let1.gr_context = None