Skip to content

Commit

Permalink
Make Protocol.receive_eof idempotent.
Browse files Browse the repository at this point in the history
This removes the need for keeping track of whether you called it or not,
especially in an asyncio context where it may be called in eof_received
or in connection_lost.
  • Loading branch information
aaugustin committed Aug 7, 2024
1 parent 8c4fd9c commit 02b3338
Show file tree
Hide file tree
Showing 2 changed files with 5 additions and 8 deletions.
5 changes: 3 additions & 2 deletions src/websockets/protocol.py
Original file line number Diff line number Diff line change
Expand Up @@ -270,10 +270,11 @@ def receive_eof(self) -> None:
- You aren't expected to call :meth:`events_received`; it won't return
any new events.
Raises:
EOFError: If :meth:`receive_eof` was called earlier.
:meth:`receive_eof` is idempotent.
"""
if self.reader.eof:
return
self.reader.feed_eof()
next(self.parser)

Expand Down
8 changes: 2 additions & 6 deletions tests/test_protocol.py
Original file line number Diff line number Diff line change
Expand Up @@ -1590,18 +1590,14 @@ def test_client_receives_eof_after_eof(self):
client.receive_data(b"\x88\x00")
self.assertConnectionClosing(client)
client.receive_eof()
with self.assertRaises(EOFError) as raised:
client.receive_eof()
self.assertEqual(str(raised.exception), "stream ended")
client.receive_eof() # this is idempotent

def test_server_receives_eof_after_eof(self):
server = Protocol(SERVER)
server.receive_data(b"\x88\x80\x3c\x3c\x3c\x3c")
self.assertConnectionClosing(server)
server.receive_eof()
with self.assertRaises(EOFError) as raised:
server.receive_eof()
self.assertEqual(str(raised.exception), "stream ended")
server.receive_eof() # this is idempotent


class TCPCloseTests(ProtocolTestCase):
Expand Down

0 comments on commit 02b3338

Please sign in to comment.