|
2 | 2 | import contextvars |
3 | 3 | import decimal |
4 | 4 | import random |
| 5 | +import socket |
5 | 6 | import weakref |
6 | 7 |
|
7 | 8 | from uvloop import _testbase as tb |
8 | 9 |
|
9 | 10 |
|
| 11 | +class _Protocol(asyncio.Protocol): |
| 12 | + def __init__(self, *, loop=None): |
| 13 | + self.done = asyncio.Future(loop=loop) |
| 14 | + |
| 15 | + def connection_lost(self, exc): |
| 16 | + if exc is None: |
| 17 | + self.done.set_result(None) |
| 18 | + else: |
| 19 | + self.done.set_exception(exc) |
| 20 | + |
| 21 | + |
10 | 22 | class _ContextBaseTests: |
11 | 23 |
|
12 | 24 | def test_task_decimal_context(self): |
@@ -126,6 +138,40 @@ async def main(): |
126 | 138 | del tracked |
127 | 139 | self.assertIsNone(ref()) |
128 | 140 |
|
| 141 | + def test_create_server_protocol_factory_context(self): |
| 142 | + cvar = contextvars.ContextVar('cvar', default='outer') |
| 143 | + factory_called_future = self.loop.create_future() |
| 144 | + proto = _Protocol(loop=self.loop) |
| 145 | + |
| 146 | + def factory(): |
| 147 | + try: |
| 148 | + self.assertEqual(cvar.get(), 'inner') |
| 149 | + except Exception as e: |
| 150 | + factory_called_future.set_exception(e) |
| 151 | + else: |
| 152 | + factory_called_future.set_result(None) |
| 153 | + |
| 154 | + return proto |
| 155 | + |
| 156 | + async def test(): |
| 157 | + cvar.set('inner') |
| 158 | + port = tb.find_free_port() |
| 159 | + srv = await self.loop.create_server(factory, '127.0.0.1', port) |
| 160 | + |
| 161 | + s = socket.socket(socket.AF_INET) |
| 162 | + with s: |
| 163 | + s.setblocking(False) |
| 164 | + await self.loop.sock_connect(s, ('127.0.0.1', port)) |
| 165 | + |
| 166 | + try: |
| 167 | + await factory_called_future |
| 168 | + finally: |
| 169 | + srv.close() |
| 170 | + await proto.done |
| 171 | + await srv.wait_closed() |
| 172 | + |
| 173 | + self.loop.run_until_complete(test()) |
| 174 | + |
129 | 175 |
|
130 | 176 | class Test_UV_Context(_ContextBaseTests, tb.UVTestCase): |
131 | 177 | pass |
|
0 commit comments