A `FakeSocket` Python class that implements the `recv` method

fakesocket.py

from dataclasses import dataclass
from typing import Optional


@dataclass
class FakeSocket:
    chunks: list[bytes] = None
    buf: bytes = None
    max_chunk_size: Optional[int] = None
    log_func: Optional[callable] = None

    def recv(self, n: int):
        if self.chunks:
            return self.chunks.pop(0)
        elif self.buf:
            if self.max_chunk_size is None:
                self.max_chunk_size = len(self.buf)
            max_chunk_size = min(n, self.max_chunk_size)
            chunk_size = random.randint(1, max_chunk_size)
            chunk, self.buf = self.buf[:chunk_size], self.buf[chunk_size:]
            self.log(dict(chunk_size=chunk_size, chunk=chunk))
            return chunk
        else:
            return b''
        
    def log(self, *args, **kwargs):
        if self.log_func:
            self.log_func(*args, **kwargs)

 
def test_receiver2():
    fake_sock = FakeSocket(
        chunks=[
            b'ChatMessage\r', b'\n4', b'3\r', b'\n{', b'"pla', b'yer', b'id"',
            b': "Dave', b'", "tex', b't', b'": "H', b'ello Wo', b'rld"}', b'Pla',
            b'yerUpdate\r', b'\n39\r\n', b'{"pl', b'ayeri', b'd": "Pa', b'ula"',
            b',', b' "x": 2', b'3', b', "y": ', b'41}',
        ],
    )
    messages_stream = get_messages_stream(fake_sock)
    assert next(messages_stream) == ChatMessage('Dave', 'Hello World')
    assert next(messages_stream) == PlayerUpdate(playerid='Paula', x=23, y=41)

    for i in range(5000):
        fake_sock = FakeSocket(
            buf=b''.join([
                b'ChatMessage\r\n43\r\n{"playerid": "Dave", "text": "Hello World"}',
                b'PlayerUpdate\r\n39\r\n{"playerid": "Paula", "x": 23, "y": 41}',
            ]),
        )
        messages_stream = get_messages_stream(fake_sock)
        assert next(messages_stream) == ChatMessage('Dave', 'Hello World')
        assert next(messages_stream) == PlayerUpdate(playerid='Paula', x=23, y=41)

    print("Excellent get_messages_stream!")

    
test_receiver2()