gh-81322: support multiple separators in StreamReader.readuntil (#16429)

This commit is contained in:
Bruce Merry 2024-04-08 18:58:02 +02:00 committed by GitHub
parent 24a2bd0481
commit 775912a51d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 102 additions and 20 deletions

View file

@ -383,6 +383,10 @@ class StreamTests(test_utils.TestCase):
stream = asyncio.StreamReader(loop=self.loop)
with self.assertRaisesRegex(ValueError, 'Separator should be'):
self.loop.run_until_complete(stream.readuntil(separator=b''))
with self.assertRaisesRegex(ValueError, 'Separator should be'):
self.loop.run_until_complete(stream.readuntil(separator=[b'']))
with self.assertRaisesRegex(ValueError, 'Separator should contain'):
self.loop.run_until_complete(stream.readuntil(separator=[]))
def test_readuntil_multi_chunks(self):
stream = asyncio.StreamReader(loop=self.loop)
@ -466,6 +470,48 @@ class StreamTests(test_utils.TestCase):
self.assertEqual(b'some dataAAA', stream._buffer)
def test_readuntil_multi_separator(self):
stream = asyncio.StreamReader(loop=self.loop)
# Simple case
stream.feed_data(b'line 1\nline 2\r')
data = self.loop.run_until_complete(stream.readuntil([b'\r', b'\n']))
self.assertEqual(b'line 1\n', data)
data = self.loop.run_until_complete(stream.readuntil([b'\r', b'\n']))
self.assertEqual(b'line 2\r', data)
self.assertEqual(b'', stream._buffer)
# First end position matches, even if that's a longer match
stream.feed_data(b'ABCDEFG')
data = self.loop.run_until_complete(stream.readuntil([b'DEF', b'BCDE']))
self.assertEqual(b'ABCDE', data)
self.assertEqual(b'FG', stream._buffer)
def test_readuntil_multi_separator_limit(self):
stream = asyncio.StreamReader(loop=self.loop, limit=3)
stream.feed_data(b'some dataA')
with self.assertRaisesRegex(asyncio.LimitOverrunError,
'is found') as cm:
self.loop.run_until_complete(stream.readuntil([b'A', b'ome dataA']))
self.assertEqual(b'some dataA', stream._buffer)
def test_readuntil_multi_separator_negative_offset(self):
# If the buffer is big enough for the smallest separator (but does
# not contain it) but too small for the largest, `offset` must not
# become negative.
stream = asyncio.StreamReader(loop=self.loop)
stream.feed_data(b'data')
readuntil_task = self.loop.create_task(stream.readuntil([b'A', b'long sep']))
self.loop.call_soon(stream.feed_data, b'Z')
self.loop.call_soon(stream.feed_data, b'Aaaa')
data = self.loop.run_until_complete(readuntil_task)
self.assertEqual(b'dataZA', data)
self.assertEqual(b'aaa', stream._buffer)
def test_readexactly_zero_or_less(self):
# Read exact number of bytes (zero or less).
stream = asyncio.StreamReader(loop=self.loop)