@@ -53,37 +53,44 @@ def test_bytes_written():
5353 assert wfile .bytes_written == 3
5454
5555
56- def test_flush_preserves_data_when_raw_write_returns_none ():
56+ class _RawWriteBlockOnce :
57+ """Mock raw.write() that returns None on the first call, then writes normally."""
58+
59+ def __init__ (self ):
60+ """Initialize _RawWriteBlockOnce."""
61+ self .call_count = 0
62+ self .written = bytearray ()
63+
64+ def __call__ (self , chunk ):
65+ """Return None on first call to simulate a blocked socket write."""
66+ self .call_count += 1
67+ if self .call_count == 1 :
68+ return (
69+ None # simulates socket returning None on first blocked write
70+ )
71+ self .written .extend (chunk )
72+ return len (chunk )
73+
74+
75+ def test_flush_when_raw_write_returns_none ():
5776 """_flush_unlocked() must not treat None from raw.write() as a byte count.
5877
5978 io.RawIOBase.write() returns None when a non-blocking socket cannot accept
6079 data. del self._write_buf[:None] is equivalent to del self._write_buf[:]
6180 which silently clears the entire buffer, truncating the response without
6281 raising an exception.
6382 """
64- data = b'x' * 32768 # larger than SOCK_WRITE_BLOCKSIZE to stress the loop
83+ data = b'x' * ( makefile . SOCK_WRITE_BLOCKSIZE * 2 ) # stress the write loop
6584
6685 sock = MockSocket ()
6786 wfile = makefile .MakeFile (sock , 'w' )
6887 wfile ._write_buf .extend (data )
6988
70- written = bytearray ()
71- call_count = 0
72-
73- def mock_raw_write (b ):
74- nonlocal call_count
75- call_count += 1
76- if call_count == 1 :
77- return (
78- None # simulates socket returning None on first blocked write
79- )
80- written .extend (b )
81- return len (b )
82-
83- wfile .raw .write = mock_raw_write
89+ mock = _RawWriteBlockOnce ()
90+ wfile .raw .write = mock
8491 wfile ._flush_unlocked ()
8592
86- assert bytes (written ) == data , (
87- f'Expected { len (data )} bytes but only { len (written )} reached raw.write(): '
93+ assert bytes (mock . written ) == data , (
94+ f'Expected { len (data )} bytes but only { len (mock . written )} reached raw.write(): '
8895 'buffer was silently discarded when raw.write() returned None'
8996 )
0 commit comments