Skip to content

Commit ac657a5

Browse files
Fix test mock setup for async context managers
Use @contextlib.asynccontextmanager decorator instead of manual __aenter__/__aexit__ implementation for mock_connect functions. Fixes test failures in: - test_transport_server_task_cleanup_on_exception - test_transport_server_task_no_cleanup_on_terminated
1 parent 410be5b commit ac657a5

File tree

1 file changed

+6
-18
lines changed

1 file changed

+6
-18
lines changed

tests/server/test_session_roaming.py

Lines changed: 6 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -289,15 +289,9 @@ async def test_transport_server_task_cleanup_on_exception():
289289
mock_read_stream = AsyncMock()
290290
mock_write_stream = AsyncMock()
291291

292-
async def mock_connect(self: Any) -> Any:
293-
class MockStreams:
294-
async def __aenter__(self) -> Any:
295-
return (mock_read_stream, mock_write_stream)
296-
297-
async def __aexit__(self, *args: Any) -> None:
298-
pass
299-
300-
return MockStreams()
292+
@contextlib.asynccontextmanager
293+
async def mock_connect() -> Any:
294+
yield (mock_read_stream, mock_write_stream)
301295

302296
async with manager.run():
303297
# Manually add transport to instances
@@ -338,15 +332,9 @@ async def test_transport_server_task_no_cleanup_on_terminated():
338332
mock_read_stream = AsyncMock()
339333
mock_write_stream = AsyncMock()
340334

341-
async def mock_connect(self: Any) -> Any:
342-
class MockStreams:
343-
async def __aenter__(self) -> Any:
344-
return (mock_read_stream, mock_write_stream)
345-
346-
async def __aexit__(self, *args: Any) -> None:
347-
pass
348-
349-
return MockStreams()
335+
@contextlib.asynccontextmanager
336+
async def mock_connect() -> Any:
337+
yield (mock_read_stream, mock_write_stream)
350338

351339
async with manager.run():
352340
# Manually add transport to instances

0 commit comments

Comments
 (0)