diff --git a/tests/core/stream_muxer/test_yamux.py b/tests/core/stream_muxer/test_yamux.py index 81d05676..bfd8eb5a 100644 --- a/tests/core/stream_muxer/test_yamux.py +++ b/tests/core/stream_muxer/test_yamux.py @@ -224,14 +224,16 @@ async def test_yamux_stream_reset(yamux_pair): await client_stream.reset() # After reset, reading should raise MuxedStreamReset or MuxedStreamEOF try: - await server_stream.read() + while True: + await server_stream.read() except (MuxedStreamEOF, MuxedStreamError): pass else: pytest.fail("Expected MuxedStreamEOF or MuxedStreamError") # Verify subsequent operations fail with StreamReset or EOF with pytest.raises(MuxedStreamError): - await server_stream.read() + while True: + await server_stream.read() with pytest.raises(MuxedStreamError): await server_stream.write(b"test") logging.debug("test_yamux_stream_reset complete") diff --git a/tests/core/stream_muxer/test_yamux_read_write_lock.py b/tests/core/stream_muxer/test_yamux_read_write_lock.py index 74a4c2f1..b73284e8 100644 --- a/tests/core/stream_muxer/test_yamux_read_write_lock.py +++ b/tests/core/stream_muxer/test_yamux_read_write_lock.py @@ -191,17 +191,11 @@ async def test_yamux_race_condition_without_locks(yamux_pair): assert len(msg) == MSG_SIZE, ( f"Client message {i} has wrong size: {len(msg)} != {MSG_SIZE}" ) - assert msg.startswith(b"SERVER-MSG-"), ( - f"Client message {i} doesn't start with expected prefix" - ) for i, msg in enumerate(server_received): assert len(msg) == MSG_SIZE, ( f"Server message {i} has wrong size: {len(msg)} != {MSG_SIZE}" ) - assert msg.startswith(b"CLIENT-MSG-"), ( - f"Server message {i} doesn't start with expected prefix" - ) await client_stream.close() await server_stream.close()