From 62c636a9c024fa1c0a63c651a1595d65a580fb6e Mon Sep 17 00:00:00 2001
From: Ben Darnell <ben@bendarnell.com>
Date: Sat, 29 Dec 2018 14:06:44 -0500
Subject: [PATCH] gen_test: Verify that contextvars work as expected

---
 tornado/test/gen_test.py | 58 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 58 insertions(+)

diff --git a/tornado/test/gen_test.py b/tornado/test/gen_test.py
index 8a84e4384..659e22605 100644
--- a/tornado/test/gen_test.py
+++ b/tornado/test/gen_test.py
@@ -16,6 +16,11 @@ from tornado.web import Application, RequestHandler, HTTPError
 
 from tornado import gen
 
+try:
+    import contextvars
+except ImportError:
+    contextvars = None  # type: ignore
+
 import typing
 
 if typing.TYPE_CHECKING:
@@ -1049,5 +1054,58 @@ class RunnerGCTest(AsyncTestCase):
         self.assertEqual(result, [None, None])
 
 
+if contextvars is not None:
+    ctx_var = contextvars.ContextVar("ctx_var")  # type: contextvars.ContextVar[int]
+
+
+@unittest.skipIf(contextvars is None, "contextvars module not present")
+class ContextVarsTest(AsyncTestCase):
+    async def native_root(self, x):
+        ctx_var.set(x)
+        await self.inner(x)
+
+    @gen.coroutine
+    def gen_root(self, x):
+        ctx_var.set(x)
+        yield
+        yield self.inner(x)
+
+    async def inner(self, x):
+        self.assertEqual(ctx_var.get(), x)
+        await self.gen_inner(x)
+        self.assertEqual(ctx_var.get(), x)
+
+        # IOLoop.run_in_executor doesn't automatically copy context
+        ctx = contextvars.copy_context()
+        await self.io_loop.run_in_executor(None, lambda: ctx.run(self.thread_inner, x))
+        self.assertEqual(ctx_var.get(), x)
+
+        # Neither does asyncio's run_in_executor.
+        await asyncio.get_event_loop().run_in_executor(
+            None, lambda: ctx.run(self.thread_inner, x)
+        )
+        self.assertEqual(ctx_var.get(), x)
+
+    @gen.coroutine
+    def gen_inner(self, x):
+        self.assertEqual(ctx_var.get(), x)
+        yield
+        self.assertEqual(ctx_var.get(), x)
+
+    def thread_inner(self, x):
+        self.assertEqual(ctx_var.get(), x)
+
+    @gen_test
+    def test_propagate(self):
+        # Verify that context vars get propagated across various
+        # combinations of native and decorated coroutines.
+        yield [
+            self.native_root(1),
+            self.native_root(2),
+            self.gen_root(3),
+            self.gen_root(4),
+        ]
+
+
 if __name__ == "__main__":
     unittest.main()
-- 
2.47.2