diff --git a/python/e2e/test_streaming_fidelity_e2e.py b/python/e2e/test_streaming_fidelity_e2e.py index a644acb838..a25957df10 100644 --- a/python/e2e/test_streaming_fidelity_e2e.py +++ b/python/e2e/test_streaming_fidelity_e2e.py @@ -5,12 +5,29 @@ import pytest from copilot import CopilotClient, RuntimeConnection +from copilot.rpc import SessionsCheckInUseRequest from copilot.session import PermissionHandler -from .testharness import E2ETestContext +from .testharness import E2ETestContext, wait_for_condition pytestmark = pytest.mark.asyncio(loop_scope="module") +RESUME_LOCK_TIMEOUT = 60.0 + + +async def _wait_for_session_lock_release(ctx: E2ETestContext, session_id: str) -> None: + async def session_lock_is_released() -> bool: + result = await ctx.client.rpc.sessions.check_in_use( + SessionsCheckInUseRequest(session_ids=[session_id]) + ) + return session_id not in result.in_use + + await wait_for_condition( + session_lock_is_released, + timeout=RESUME_LOCK_TIMEOUT, + timeout_message=f"Timed out waiting for session '{session_id}' to release its lock.", + ) + class TestStreamingFidelity: async def test_should_produce_delta_events_when_streaming_is_enabled(self, ctx: E2ETestContext): @@ -70,9 +87,12 @@ async def test_should_produce_deltas_after_session_resume(self, ctx: E2ETestCont session = await ctx.client.create_session( on_permission_request=PermissionHandler.approve_all, streaming=False ) + session_id = session.session_id await session.send_and_wait("What is 3 + 6?") await session.disconnect() + await _wait_for_session_lock_release(ctx, session_id) + # Resume using a new client github_token = ( "fake-token-for-e2e-tests" if os.environ.get("GITHUB_ACTIONS") == "true" else None @@ -86,7 +106,7 @@ async def test_should_produce_deltas_after_session_resume(self, ctx: E2ETestCont try: session2 = await new_client.resume_session( - session.session_id, + session_id, on_permission_request=PermissionHandler.approve_all, streaming=True, ) @@ -126,6 +146,8 @@ async def test_should_not_produce_deltas_after_session_resume_with_streaming_dis session_id = session.session_id await session.disconnect() + await _wait_for_session_lock_release(ctx, session_id) + # Resume with streaming disabled new_client = CopilotClient( connection=RuntimeConnection.for_stdio(path=ctx.cli_path),