feat(dirty): update arbiter for binary protocol

Update arbiter tests to work with the binary protocol:
- Update MockStreamWriter to decode binary messages
- Import binary protocol constants from module level
This commit is contained in:
Benoit Chesneau 2026-02-11 23:03:40 +01:00
parent 6d2139bb6c
commit 98b1b649c2

View File

@ -14,7 +14,12 @@ import pytest
from gunicorn.config import Config from gunicorn.config import Config
from gunicorn.dirty.arbiter import DirtyArbiter from gunicorn.dirty.arbiter import DirtyArbiter
from gunicorn.dirty.errors import DirtyError from gunicorn.dirty.errors import DirtyError
from gunicorn.dirty.protocol import DirtyProtocol, make_request from gunicorn.dirty.protocol import (
DirtyProtocol,
BinaryProtocol,
make_request,
HEADER_SIZE,
)
class MockStreamWriter: class MockStreamWriter:
@ -29,16 +34,22 @@ class MockStreamWriter:
self._buffer += data self._buffer += data
async def drain(self): async def drain(self):
while len(self._buffer) >= DirtyProtocol.HEADER_SIZE: # Decode the buffer to extract messages using binary protocol
length = struct.unpack( while len(self._buffer) >= HEADER_SIZE:
DirtyProtocol.HEADER_FORMAT, # Decode header to get payload length
self._buffer[:DirtyProtocol.HEADER_SIZE] _, _, length = BinaryProtocol.decode_header(
)[0] self._buffer[:HEADER_SIZE]
total_size = DirtyProtocol.HEADER_SIZE + length )
total_size = HEADER_SIZE + length
if len(self._buffer) >= total_size: if len(self._buffer) >= total_size:
msg_data = self._buffer[DirtyProtocol.HEADER_SIZE:total_size] msg_data = self._buffer[:total_size]
self._buffer = self._buffer[total_size:] self._buffer = self._buffer[total_size:]
self.messages.append(DirtyProtocol.decode(msg_data)) # decode_message returns (msg_type_str, request_id, payload_dict)
msg_type_str, request_id, payload_dict = BinaryProtocol.decode_message(msg_data)
# Reconstruct the dict format for backwards compatibility
result = {"type": msg_type_str, "id": request_id}
result.update(payload_dict)
self.messages.append(result)
else: else:
break break