test: add tests for tunnel reconnection retry logic

Add unit tests covering the new retry logic in handle_async (lease.py)
and the channel_ready timeout in connect_router_stream (streams.py) to
satisfy the 80% diff coverage requirement.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: Ambient Code Bot

python/packages/jumpstarter/jumpstarter/client/lease_test.py

@@ -1,10 +1,13 @@
 import asyncio
 import logging
 import sys
+from contextlib import asynccontextmanager
 from datetime import datetime, timedelta, timezone
 from unittest.mock import AsyncMock, Mock, patch
 
+import grpc
 import pytest
+from grpc.aio import AioRpcError
 from rich.console import Console
 
 from jumpstarter.client.exceptions import LeaseError
@@ -554,3 +557,270 @@ async def get_then_fail():
         callback.assert_called()
         _, remain_arg = callback.call_args[0]
         assert remain_arg == timedelta(0)
+
+
+def _make_aio_rpc_error(code, details="error"):
+    """Helper to construct an AioRpcError."""
+    return AioRpcError(
+        code=code,
+        initial_metadata=grpc.aio.Metadata(),
+        trailing_metadata=grpc.aio.Metadata(),
+        details=details,
+        debug_error_string=None,
+    )
+
+
+def _make_lease_for_handle():
+    """Create a minimal Lease for testing handle_async."""
+    lease = object.__new__(Lease)
+    lease.name = "test-lease"
+    lease.dial_timeout = 5.0
+    lease.tls_config = Mock()
+    lease.grpc_options = {}
+    lease.controller = Mock()
+    lease.lease_transferred = False
+    return lease
+
+
+class TestHandleAsyncTransientDialRetry:
+    """Tests for transient gRPC error retry in handle_async Dial phase."""
+
+    @pytest.mark.anyio
+    async def test_dial_retries_on_unavailable_then_succeeds(self):
+        """Dial should retry on UNAVAILABLE and succeed on the next attempt."""
+        lease = _make_lease_for_handle()
+
+        dial_response = Mock(router_endpoint="ep", router_token="tok")
+        call_count = 0
+
+        async def dial_side_effect(req):
+            nonlocal call_count
+            call_count += 1
+            if call_count == 1:
+                raise _make_aio_rpc_error(grpc.StatusCode.UNAVAILABLE, "tunnel dropped")
+            return dial_response
+
+        lease.controller.Dial = AsyncMock(side_effect=dial_side_effect)
+
+        with patch("jumpstarter.client.lease.sleep", new_callable=AsyncMock):
+            with patch("jumpstarter.client.lease.connect_router_stream") as mock_router:
+
+                @asynccontextmanager
+                async def fake_router(*args, **kwargs):
+                    yield
+
+                mock_router.side_effect = fake_router
+                await lease.handle_async(Mock())
+
+        assert call_count == 2
+
+    @pytest.mark.anyio
+    async def test_dial_transient_error_returns_after_timeout(self):
+        """Dial should give up and return when dial_timeout is exceeded."""
+        lease = _make_lease_for_handle()
+        lease.dial_timeout = 0.0  # already expired
+
+        lease.controller.Dial = AsyncMock(
+            side_effect=_make_aio_rpc_error(grpc.StatusCode.UNAVAILABLE, "tunnel dropped"),
+        )
+
+        with patch("jumpstarter.client.lease.sleep", new_callable=AsyncMock):
+            await lease.handle_async(Mock())
+
+        # Should return without raising
+        lease.controller.Dial.assert_called_once()
+
+    @pytest.mark.anyio
+    @pytest.mark.parametrize(
+        "status_code",
+        [grpc.StatusCode.RESOURCE_EXHAUSTED, grpc.StatusCode.ABORTED, grpc.StatusCode.INTERNAL],
+        ids=["RESOURCE_EXHAUSTED", "ABORTED", "INTERNAL"],
+    )
+    async def test_dial_retries_multiple_transient_codes(self, status_code):
+        """Dial should retry on RESOURCE_EXHAUSTED, ABORTED, INTERNAL."""
+        lease = _make_lease_for_handle()
+        dial_response = Mock(router_endpoint="ep", router_token="tok")
+        call_count = 0
+
+        async def dial_side_effect(req):
+            nonlocal call_count
+            call_count += 1
+            if call_count == 1:
+                raise _make_aio_rpc_error(status_code, "transient")
+            return dial_response
+
+        lease.controller.Dial = AsyncMock(side_effect=dial_side_effect)
+
+        with patch("jumpstarter.client.lease.sleep", new_callable=AsyncMock):
+            with patch("jumpstarter.client.lease.connect_router_stream") as mock_router:
+
+                @asynccontextmanager
+                async def fake_router(*args, **kwargs):
+                    yield
+
+                mock_router.side_effect = fake_router
+                await lease.handle_async(Mock())
+
+        assert call_count == 2, f"Expected 2 calls for {status_code}, got {call_count}"
+
+
+class TestHandleAsyncRouterRetry:
+    """Tests for router connection retry in handle_async."""
+
+    @pytest.mark.anyio
+    async def test_router_retries_on_transient_error_then_succeeds(self):
+        """Router connection should retry on transient error, re-dial, then succeed."""
+        lease = _make_lease_for_handle()
+        dial_response = Mock(router_endpoint="ep", router_token="tok")
+        lease.controller.Dial = AsyncMock(return_value=dial_response)
+
+        connect_count = 0
+
+        @asynccontextmanager
+        async def fake_router(*args, **kwargs):
+            nonlocal connect_count
+            connect_count += 1
+            if connect_count == 1:
+                raise _make_aio_rpc_error(grpc.StatusCode.UNAVAILABLE, "router unreachable")
+            yield
+
+        with patch("jumpstarter.client.lease.sleep", new_callable=AsyncMock):
+            with patch("jumpstarter.client.lease.connect_router_stream", side_effect=fake_router):
+                await lease.handle_async(Mock())
+
+        assert connect_count == 2
+        # Dial called once for initial + once for re-dial
+        assert lease.controller.Dial.call_count == 2
+
+    @pytest.mark.anyio
+    async def test_router_non_transient_error_returns_immediately(self):
+        """Router connection should not retry on non-transient errors."""
+        lease = _make_lease_for_handle()
+        dial_response = Mock(router_endpoint="ep", router_token="tok")
+        lease.controller.Dial = AsyncMock(return_value=dial_response)
+
+        @asynccontextmanager
+        async def fail_router(*args, **kwargs):
+            raise _make_aio_rpc_error(grpc.StatusCode.PERMISSION_DENIED, "no access")
+            yield  # pragma: no cover
+
+        with patch("jumpstarter.client.lease.connect_router_stream", side_effect=fail_router):
+            await lease.handle_async(Mock())
+
+        # Only the initial Dial, no re-dial
+        assert lease.controller.Dial.call_count == 1
+
+    @pytest.mark.anyio
+    async def test_router_transient_error_returns_after_timeout(self):
+        """Router should give up when dial_timeout is exceeded."""
+        lease = _make_lease_for_handle()
+        lease.dial_timeout = 0.0  # already expired
+        dial_response = Mock(router_endpoint="ep", router_token="tok")
+        lease.controller.Dial = AsyncMock(return_value=dial_response)
+
+        @asynccontextmanager
+        async def fail_router(*args, **kwargs):
+            raise _make_aio_rpc_error(grpc.StatusCode.UNAVAILABLE, "unreachable")
+            yield  # pragma: no cover
+
+        with patch("jumpstarter.client.lease.connect_router_stream", side_effect=fail_router):
+            await lease.handle_async(Mock())
+
+        # Only one Dial (initial), no retry
+        assert lease.controller.Dial.call_count == 1
+
+    @pytest.mark.anyio
+    async def test_router_redial_failure_is_swallowed(self):
+        """When re-dial fails during router retry, the error is logged and retry continues."""
+        lease = _make_lease_for_handle()
+        dial_response = Mock(router_endpoint="ep", router_token="tok")
+
+        dial_count = 0
+
+        async def dial_side_effect(req):
+            nonlocal dial_count
+            dial_count += 1
+            if dial_count == 1:
+                return dial_response
+            if dial_count == 2:
+                # Re-dial fails
+                raise _make_aio_rpc_error(grpc.StatusCode.UNAVAILABLE, "re-dial failed")
+            return dial_response
+
+        lease.controller.Dial = AsyncMock(side_effect=dial_side_effect)
+
+        connect_count = 0
+
+        @asynccontextmanager
+        async def fake_router(*args, **kwargs):
+            nonlocal connect_count
+            connect_count += 1
+            if connect_count <= 2:
+                raise _make_aio_rpc_error(grpc.StatusCode.UNAVAILABLE, "router fail")
+            yield
+
+        with patch("jumpstarter.client.lease.sleep", new_callable=AsyncMock):
+            with patch("jumpstarter.client.lease.connect_router_stream", side_effect=fake_router):
+                await lease.handle_async(Mock())
+
+        # Should have retried: connect fails, re-dial fails, connect fails again,
+        # re-dial succeeds, third connect succeeds
+        assert connect_count == 3
+        assert dial_count == 3
+
+    @pytest.mark.anyio
+    async def test_router_oserror_retries_then_succeeds(self):
+        """Router connection should retry on OSError, then succeed."""
+        lease = _make_lease_for_handle()
+        dial_response = Mock(router_endpoint="ep", router_token="tok")
+        lease.controller.Dial = AsyncMock(return_value=dial_response)
+
+        connect_count = 0
+
+        @asynccontextmanager
+        async def fake_router(*args, **kwargs):
+            nonlocal connect_count
+            connect_count += 1
+            if connect_count == 1:
+                raise OSError("Connection refused")
+            yield
+
+        with patch("jumpstarter.client.lease.sleep", new_callable=AsyncMock):
+            with patch("jumpstarter.client.lease.connect_router_stream", side_effect=fake_router):
+                await lease.handle_async(Mock())
+
+        assert connect_count == 2
+
+    @pytest.mark.anyio
+    async def test_router_oserror_returns_after_timeout(self):
+        """Router should give up on OSError when dial_timeout is exceeded."""
+        lease = _make_lease_for_handle()
+        lease.dial_timeout = 0.0  # already expired
+        dial_response = Mock(router_endpoint="ep", router_token="tok")
+        lease.controller.Dial = AsyncMock(return_value=dial_response)
+
+        @asynccontextmanager
+        async def fail_router(*args, **kwargs):
+            raise OSError("Connection refused")
+            yield  # pragma: no cover
+
+        with patch("jumpstarter.client.lease.connect_router_stream", side_effect=fail_router):
+            await lease.handle_async(Mock())
+
+        # Only the initial Dial, no retry
+        assert lease.controller.Dial.call_count == 1
+
+
+class TestTransientGrpcCodes:
+    """Tests for the _TRANSIENT_GRPC_CODES class attribute."""
+
+    def test_contains_expected_codes(self):
+        assert grpc.StatusCode.UNAVAILABLE in Lease._TRANSIENT_GRPC_CODES
+        assert grpc.StatusCode.RESOURCE_EXHAUSTED in Lease._TRANSIENT_GRPC_CODES
+        assert grpc.StatusCode.ABORTED in Lease._TRANSIENT_GRPC_CODES
+        assert grpc.StatusCode.INTERNAL in Lease._TRANSIENT_GRPC_CODES
+
+    def test_does_not_contain_non_transient_codes(self):
+        assert grpc.StatusCode.PERMISSION_DENIED not in Lease._TRANSIENT_GRPC_CODES
+        assert grpc.StatusCode.NOT_FOUND not in Lease._TRANSIENT_GRPC_CODES
+        assert grpc.StatusCode.FAILED_PRECONDITION not in Lease._TRANSIENT_GRPC_CODES

python/packages/jumpstarter/jumpstarter/common/streams_test.py

@@ -0,0 +1,88 @@
+import asyncio
+from contextlib import asynccontextmanager
+from unittest.mock import AsyncMock, Mock, patch
+
+import grpc
+import pytest
+from grpc.aio import AioRpcError
+
+from jumpstarter.common.streams import connect_router_stream
+
+
+class TestConnectRouterStreamChannelReady:
+    """Tests for the channel_ready timeout logic in connect_router_stream."""
+
+    @pytest.mark.anyio
+    async def test_raises_unavailable_on_channel_ready_timeout(self):
+        """When channel_ready() times out, an AioRpcError with UNAVAILABLE should be raised."""
+        mock_channel = Mock()
+
+        # Make channel_ready() return a coroutine that never completes
+        async def hang_forever():
+            await asyncio.sleep(999)
+
+        mock_channel.channel_ready = Mock(return_value=hang_forever())
+
+        @asynccontextmanager
+        async def fake_secure_channel(*args, **kwargs):
+            yield mock_channel
+
+        with (
+            patch("jumpstarter.common.streams.ssl_channel_credentials", new_callable=AsyncMock),
+            patch("jumpstarter.common.streams.aio_secure_channel", side_effect=fake_secure_channel),
+            patch("grpc.composite_channel_credentials", return_value=Mock()),
+            patch("grpc.access_token_call_credentials", return_value=Mock()),
+        ):
+            with pytest.raises(AioRpcError) as exc_info:
+                async with connect_router_stream(
+                    "endpoint:443", "token", Mock(), Mock(), {}, channel_ready_timeout=0.01
+                ):
+                    pass  # pragma: no cover
+
+            assert exc_info.value.code() == grpc.StatusCode.UNAVAILABLE
+            assert "Timed out" in str(exc_info.value.details())
+
+    @pytest.mark.anyio
+    async def test_proceeds_when_channel_ready_succeeds(self):
+        """When channel_ready() succeeds quickly, the stream should be set up normally."""
+        mock_channel = Mock()
+
+        # channel_ready() resolves immediately
+        async def ready_immediately():
+            pass
+
+        mock_channel.channel_ready = Mock(return_value=ready_immediately())
+
+        mock_context = Mock()
+
+        @asynccontextmanager
+        async def fake_secure_channel(*args, **kwargs):
+            yield mock_channel
+
+        @asynccontextmanager
+        async def fake_router_stream(*args, **kwargs):
+            yield Mock()
+
+        @asynccontextmanager
+        async def fake_forward(*args, **kwargs):
+            yield
+
+        with (
+            patch("jumpstarter.common.streams.ssl_channel_credentials", new_callable=AsyncMock),
+            patch("jumpstarter.common.streams.aio_secure_channel", side_effect=fake_secure_channel),
+            patch("grpc.composite_channel_credentials", return_value=Mock()),
+            patch("grpc.access_token_call_credentials", return_value=Mock()),
+            patch("jumpstarter.common.streams.router_pb2_grpc.RouterServiceStub") as mock_stub_cls,
+            patch("jumpstarter.common.streams.RouterStream", side_effect=fake_router_stream),
+            patch("jumpstarter.common.streams.forward_stream", side_effect=fake_forward),
+        ):
+            mock_stub = Mock()
+            mock_stub.Stream.return_value = mock_context
+            mock_stub_cls.return_value = mock_stub
+
+            async with connect_router_stream(
+                "endpoint:443", "token", Mock(), Mock(), {}, channel_ready_timeout=5
+            ):
+                pass  # Successfully entered the context
+
+            mock_channel.channel_ready.assert_called_once()