|
1 | 1 | import asyncio
|
2 | 2 | import sys
|
3 |
| -from typing import Any, NoReturn, Optional |
| 3 | +from typing import Any, List, NoReturn, Optional |
4 | 4 | from unittest import mock
|
5 | 5 |
|
6 | 6 | import pytest
|
@@ -1203,3 +1203,71 @@ async def test_ws_connect_with_wrong_ssl_type(aiohttp_client: AiohttpClient) ->
|
1203 | 1203 |
|
1204 | 1204 | with pytest.raises(TypeError, match="ssl should be SSLContext, .*"):
|
1205 | 1205 | await session.ws_connect("/", ssl=42)
|
| 1206 | + |
| 1207 | + |
| 1208 | +async def test_websocket_connection_not_closed_properly( |
| 1209 | + aiohttp_client: AiohttpClient, |
| 1210 | +) -> None: |
| 1211 | + """Test that closing the connection via __del__ does not raise an exception.""" |
| 1212 | + |
| 1213 | + async def handler(request: web.Request) -> NoReturn: |
| 1214 | + ws = web.WebSocketResponse() |
| 1215 | + await ws.prepare(request) |
| 1216 | + await ws.close() |
| 1217 | + assert False |
| 1218 | + |
| 1219 | + app = web.Application() |
| 1220 | + app.router.add_route("GET", "/", handler) |
| 1221 | + client = await aiohttp_client(app) |
| 1222 | + resp = await client.ws_connect("/") |
| 1223 | + assert resp._conn is not None |
| 1224 | + # Simulate the connection not being closed properly |
| 1225 | + # https://github.com/aio-libs/aiohttp/issues/9880 |
| 1226 | + resp._conn.release() |
| 1227 | + |
| 1228 | + # Clean up so the test does not leak |
| 1229 | + await resp.close() |
| 1230 | + |
| 1231 | + |
| 1232 | +async def test_websocket_connection_cancellation( |
| 1233 | + aiohttp_client: AiohttpClient, loop: asyncio.AbstractEventLoop |
| 1234 | +) -> None: |
| 1235 | + """Test canceling the WebSocket connection task does not raise an exception in __del__.""" |
| 1236 | + |
| 1237 | + async def handler(request: web.Request) -> NoReturn: |
| 1238 | + ws = web.WebSocketResponse() |
| 1239 | + await ws.prepare(request) |
| 1240 | + await ws.close() |
| 1241 | + assert False |
| 1242 | + |
| 1243 | + app = web.Application() |
| 1244 | + app.router.add_route("GET", "/", handler) |
| 1245 | + |
| 1246 | + sync_future: "asyncio.Future[List[aiohttp.ClientWebSocketResponse]]" = ( |
| 1247 | + loop.create_future() |
| 1248 | + ) |
| 1249 | + client = await aiohttp_client(app) |
| 1250 | + |
| 1251 | + async def websocket_task() -> None: |
| 1252 | + resp = await client.ws_connect("/") |
| 1253 | + assert resp is not None # ensure we hold a reference to the response |
| 1254 | + # The test harness will cleanup the unclosed websocket |
| 1255 | + # for us, so we need to copy the websockets to ensure |
| 1256 | + # we can control the cleanup |
| 1257 | + sync_future.set_result(client._websockets.copy()) |
| 1258 | + client._websockets.clear() |
| 1259 | + await asyncio.sleep(0) |
| 1260 | + |
| 1261 | + task = loop.create_task(websocket_task()) |
| 1262 | + websockets = await sync_future |
| 1263 | + task.cancel() |
| 1264 | + with pytest.raises(asyncio.CancelledError): |
| 1265 | + await task |
| 1266 | + |
| 1267 | + websocket = websockets.pop() |
| 1268 | + # Call the `__del__` methods manually since when it gets gc'd it not reproducible |
| 1269 | + del websocket._response |
| 1270 | + |
| 1271 | + # Cleanup properly |
| 1272 | + websocket._response = mock.Mock() |
| 1273 | + await websocket.close() |
0 commit comments