Skip to content

Commit

Permalink
added a break condition to _ctrl_read_forever
Browse files Browse the repository at this point in the history
  • Loading branch information
evalott100 committed Apr 8, 2024
1 parent 8d95812 commit e19d68c
Show file tree
Hide file tree
Showing 2 changed files with 39 additions and 0 deletions.
8 changes: 8 additions & 0 deletions src/pandablocks/asyncio.py
Original file line number Diff line number Diff line change
Expand Up @@ -99,8 +99,16 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
await self.close()

async def _ctrl_read_forever(self, reader: asyncio.StreamReader):
"""Continually read data from the stream reader and add to the data queue.
Args:
reader: The `StreamReader` to read from
"""
while True:
received = await reader.read(4096)
if received == b"":
raise ConnectionError("Received an empty packet. Closing connection.")

try:
to_send = self._ctrl_connection.receive_bytes(received)
await self._ctrl_stream.write_and_drain(to_send)
Expand Down
31 changes: 31 additions & 0 deletions tests/test_asyncio.py
Original file line number Diff line number Diff line change
Expand Up @@ -104,6 +104,37 @@ async def test_asyncio_data_timeout(dummy_server_async, fast_dump):
"This goes forever, when it runs out of data we will get our timeout"


async def test_asyncio_empty_frame_error():
dummy_data = [b"ABC"] * 10 + [b""]
dummy_data_iter = iter(dummy_data)

async def dummy_read(n):
return dummy_data_iter.__next__()

reader = asyncio.StreamReader()
reader.read = dummy_read

written = []

class DummyControlStream:
async def write_and_drain(self, data):
written.append(data)

class DummyControlConnection:
def receive_bytes(self, data):
return data

client = AsyncioClient("localhost")
client._ctrl_stream = DummyControlStream()
client._ctrl_connection = DummyControlConnection()
with pytest.raises(
ConnectionError, match="Received an empty packet. Closing connection."
):
await client._ctrl_read_forever(reader)
assert written == dummy_data[:-1]
client.close()


@pytest.mark.asyncio
async def test_asyncio_connects(dummy_server_async: DummyServer):
async with AsyncioClient("localhost") as client:
Expand Down

0 comments on commit e19d68c

Please sign in to comment.