Skip to content

Commit b21ae98

Browse files
[PR #10971/1ee187c0 backport][3.12] Fix WebSocketResponse.prepared not correctly reflect the WebSocket's prepared state (#10983)
Co-authored-by: J. Nick Koston <[email protected]> fixes #6009
1 parent bfe0bd1 commit b21ae98

File tree

4 files changed

+119
-1
lines changed

4 files changed

+119
-1
lines changed

CHANGES/6009.bugfix.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Fixed ``WebSocketResponse.prepared`` property to correctly reflect the prepared state, especially during timeout scenarios -- by :user:`bdraco`

aiohttp/web_ws.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -354,6 +354,10 @@ def can_prepare(self, request: BaseRequest) -> WebSocketReady:
354354
else:
355355
return WebSocketReady(True, protocol)
356356

357+
@property
358+
def prepared(self) -> bool:
359+
return self._writer is not None
360+
357361
@property
358362
def closed(self) -> bool:
359363
return self._closed

tests/test_web_websocket.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -639,4 +639,4 @@ async def test_get_extra_info(
639639
await ws.prepare(req)
640640
ws._writer = ws_transport
641641

642-
assert ws.get_extra_info(valid_key, default_value) == expected_result
642+
assert expected_result == ws.get_extra_info(valid_key, default_value)

tests/test_web_websocket_functional.py

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1281,3 +1281,116 @@ async def handler(request: web.Request) -> web.WebSocketResponse:
12811281
)
12821282
await client.server.close()
12831283
assert close_code == WSCloseCode.OK
1284+
1285+
1286+
async def test_websocket_prepare_timeout_close_issue(
1287+
loop: asyncio.AbstractEventLoop, aiohttp_client: AiohttpClient
1288+
) -> None:
1289+
"""Test that WebSocket can handle prepare with early returns.
1290+
1291+
This is a regression test for issue #6009 where the prepared property
1292+
incorrectly checked _payload_writer instead of _writer.
1293+
"""
1294+
1295+
async def handler(request: web.Request) -> web.WebSocketResponse:
1296+
ws = web.WebSocketResponse()
1297+
assert ws.can_prepare(request)
1298+
await ws.prepare(request)
1299+
await ws.send_str("test")
1300+
await ws.close()
1301+
return ws
1302+
1303+
app = web.Application()
1304+
app.router.add_route("GET", "/ws", handler)
1305+
client = await aiohttp_client(app)
1306+
1307+
# Connect via websocket
1308+
ws = await client.ws_connect("/ws")
1309+
msg = await ws.receive()
1310+
assert msg.type is WSMsgType.TEXT
1311+
assert msg.data == "test"
1312+
await ws.close()
1313+
1314+
1315+
async def test_websocket_prepare_timeout_from_issue_reproducer(
1316+
loop: asyncio.AbstractEventLoop, aiohttp_client: AiohttpClient
1317+
) -> None:
1318+
"""Test websocket behavior when prepare is interrupted.
1319+
1320+
This test verifies the fix for issue #6009 where close() would
1321+
fail after prepare() was interrupted.
1322+
"""
1323+
prepare_complete = asyncio.Event()
1324+
close_complete = asyncio.Event()
1325+
1326+
async def handler(request: web.Request) -> web.WebSocketResponse:
1327+
ws = web.WebSocketResponse()
1328+
1329+
# Prepare the websocket
1330+
await ws.prepare(request)
1331+
prepare_complete.set()
1332+
1333+
# Send a message to confirm connection works
1334+
await ws.send_str("connected")
1335+
1336+
# Wait for client to close
1337+
msg = await ws.receive()
1338+
assert msg.type is WSMsgType.CLOSE
1339+
await ws.close()
1340+
close_complete.set()
1341+
1342+
return ws
1343+
1344+
app = web.Application()
1345+
app.router.add_route("GET", "/ws", handler)
1346+
client = await aiohttp_client(app)
1347+
1348+
# Connect and verify the connection works
1349+
ws = await client.ws_connect("/ws")
1350+
await prepare_complete.wait()
1351+
1352+
msg = await ws.receive()
1353+
assert msg.type is WSMsgType.TEXT
1354+
assert msg.data == "connected"
1355+
1356+
# Close the connection
1357+
await ws.close()
1358+
await close_complete.wait()
1359+
1360+
1361+
async def test_websocket_prepared_property(
1362+
loop: asyncio.AbstractEventLoop, aiohttp_client: AiohttpClient
1363+
) -> None:
1364+
"""Test that WebSocketResponse.prepared property correctly reflects state."""
1365+
prepare_called = asyncio.Event()
1366+
1367+
async def handler(request: web.Request) -> web.WebSocketResponse:
1368+
ws = web.WebSocketResponse()
1369+
1370+
# Initially not prepared
1371+
initial_state = ws.prepared
1372+
assert not initial_state
1373+
1374+
# After prepare() is called, should be prepared
1375+
await ws.prepare(request)
1376+
prepare_called.set()
1377+
1378+
# Check prepared state
1379+
prepared_state = ws.prepared
1380+
assert prepared_state
1381+
1382+
# Send a message to verify the connection works
1383+
await ws.send_str("test")
1384+
await ws.close()
1385+
return ws
1386+
1387+
app = web.Application()
1388+
app.router.add_route("GET", "/", handler)
1389+
client = await aiohttp_client(app)
1390+
1391+
ws = await client.ws_connect("/")
1392+
await prepare_called.wait()
1393+
msg = await ws.receive()
1394+
assert msg.type is WSMsgType.TEXT
1395+
assert msg.data == "test"
1396+
await ws.close()

0 commit comments

Comments
 (0)