Skip to content

Commit 57f3ae9

Browse files
committed
fix fragmentation edge case on empty payload
1 parent 0eed515 commit 57f3ae9

File tree

3 files changed

+25
-2
lines changed

3 files changed

+25
-2
lines changed

rsocket/frame_fragmenter.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,10 @@ def _get_next_fragment_body_size(self) -> int:
3838

3939
def __iter__(self):
4040

41+
if self._data_length == 0 and self._metadata_length == 0:
42+
yield Fragment(None, None, is_last=True, is_first=True)
43+
return
44+
4145
data_reader = BytesIO(self.data)
4246
metadata_reader = BytesIO(self.metadata)
4347

tests/rsocket/test_fragments.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ def test_create_data():
1515

1616

1717
@pytest.mark.parametrize('data, metadata, fragment_size_bytes, expected_frame_count', (
18+
(b'', b'', 64, 1), # empty payload
1819
(b'', create_data(b'123abc456def', 20), 64, 5), # only data
1920
(create_data(b'123abc456def', 20), b'', 64, 5), # only metadata
2021
(create_data(b'123abc456def', 20), create_data(b'123abc456def', 20, 55), 64, 6), # metadata fits in first frame

tests/rsocket/test_request_response.py

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -112,15 +112,33 @@ async def request_response(self, payload: Payload):
112112
assert response.metadata == b'(server (client metadata))'
113113

114114

115-
async def test_request_response_fragmented(lazy_pipe):
115+
@pytest.mark.parametrize('data_size_multiplier', (
116+
1,
117+
10
118+
))
119+
async def test_request_response_fragmented(lazy_pipe, data_size_multiplier):
116120
class Handler(BaseRequestHandler):
117121
async def request_response(self, request: Payload):
118122
return future_from_payload(request)
119123

120124
async with lazy_pipe(
121125
server_arguments={'handler_factory': Handler, 'fragment_size_bytes': 64},
122126
client_arguments={'fragment_size_bytes': 64}) as (server, client):
123-
data = b'dog-dog-dog-dog-dog-dog-dog-dog-dog' * 10
127+
data = b'dog-dog-dog-dog-dog-dog-dog-dog-dog' * data_size_multiplier
124128
response = await client.request_response(Payload(data, b'cat'))
125129

126130
assert response.data == b'data: ' + data
131+
132+
133+
async def test_request_response_fragmented_empty_payload(lazy_pipe):
134+
class Handler(BaseRequestHandler):
135+
async def request_response(self, request: Payload):
136+
return create_future(Payload())
137+
138+
async with lazy_pipe(
139+
server_arguments={'handler_factory': Handler, 'fragment_size_bytes': 64},
140+
client_arguments={'fragment_size_bytes': 64}) as (server, client):
141+
response = await client.request_response(Payload())
142+
143+
assert response.data == b''
144+
assert response.metadata == b''

0 commit comments

Comments
 (0)