open-swe/tests/test_github_proxy_refresh.py

"""Tests for mid-run GitHub proxy token refresh."""

from __future__ import annotations

from datetime import UTC, datetime, timedelta
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from agent.utils import github_proxy
from agent.utils.github_proxy import (
    PROXY_TOKEN_FALLBACK_TTL,
    clear_proxy_token_expiry,
    maybe_refresh_proxy_token,
    proxy_token_needs_refresh,
    record_proxy_token_expiry,
)


@pytest.fixture(autouse=True)
def _clear_state() -> None:
    github_proxy._PROXY_TOKEN_EXPIRY.clear()
    yield
    github_proxy._PROXY_TOKEN_EXPIRY.clear()


class TestProxyTokenNeedsRefresh:
    def test_false_when_no_record(self) -> None:
        assert proxy_token_needs_refresh("thread-1") is False

    def test_false_when_thread_id_missing(self) -> None:
        assert proxy_token_needs_refresh(None) is False

    def test_true_when_near_expiry(self) -> None:
        now = datetime(2025, 1, 1, 12, 0, 0, tzinfo=UTC)
        record_proxy_token_expiry("thread-1", now + timedelta(minutes=2))
        assert proxy_token_needs_refresh("thread-1", now=now) is True

    def test_false_when_far_from_expiry(self) -> None:
        now = datetime(2025, 1, 1, 12, 0, 0, tzinfo=UTC)
        record_proxy_token_expiry("thread-1", now + timedelta(minutes=55))
        assert proxy_token_needs_refresh("thread-1", now=now) is False

    def test_parses_iso_z_suffix(self) -> None:
        now = datetime(2025, 1, 1, 12, 0, 0, tzinfo=UTC)
        record_proxy_token_expiry("thread-1", "2025-01-01T12:03:00Z")
        assert proxy_token_needs_refresh("thread-1", now=now) is True

    def test_fallback_ttl_when_expiry_unknown(self) -> None:
        now = datetime(2025, 1, 1, 12, 0, 0, tzinfo=UTC)
        record_proxy_token_expiry("thread-1", None)
        github_proxy._PROXY_TOKEN_EXPIRY["thread-1"] = (None, now, None)
        assert proxy_token_needs_refresh("thread-1", now=now) is False
        later = now + PROXY_TOKEN_FALLBACK_TTL
        assert proxy_token_needs_refresh("thread-1", now=later) is True

    def test_clear_removes_record(self) -> None:
        record_proxy_token_expiry("thread-1", datetime.now(UTC))
        clear_proxy_token_expiry("thread-1")
        assert proxy_token_needs_refresh("thread-1") is False


class TestMaybeRefreshProxyToken:
    @pytest.mark.asyncio
    async def test_skips_when_not_langsmith(self) -> None:
        now = datetime(2025, 1, 1, 12, 0, 0, tzinfo=UTC)
        record_proxy_token_expiry("thread-1", now + timedelta(minutes=1))
        with patch.dict("os.environ", {"SANDBOX_TYPE": "local"}):
            assert await maybe_refresh_proxy_token("thread-1", now=now) is False

    @pytest.mark.asyncio
    async def test_skips_when_not_near_expiry(self) -> None:
        now = datetime(2025, 1, 1, 12, 0, 0, tzinfo=UTC)
        record_proxy_token_expiry("thread-1", now + timedelta(minutes=55))
        with patch.dict("os.environ", {"SANDBOX_TYPE": "langsmith"}):
            assert await maybe_refresh_proxy_token("thread-1", now=now) is False

    @pytest.mark.asyncio
    async def test_skips_when_no_sandbox(self) -> None:
        now = datetime(2025, 1, 1, 12, 0, 0, tzinfo=UTC)
        record_proxy_token_expiry("thread-1", now + timedelta(minutes=1))
        with (
            patch.dict("os.environ", {"SANDBOX_TYPE": "langsmith"}),
            patch.dict(github_proxy.SANDBOX_BACKENDS, {}, clear=True),
        ):
            assert await maybe_refresh_proxy_token("thread-1", now=now) is False

    @pytest.mark.asyncio
    async def test_refreshes_and_records_new_expiry(self) -> None:
        now = datetime(2025, 1, 1, 12, 0, 0, tzinfo=UTC)
        record_proxy_token_expiry("thread-1", now + timedelta(minutes=1))
        backend = MagicMock(id="sb-1")
        new_expiry = "2025-01-01T13:00:00Z"

        with (
            patch.dict("os.environ", {"SANDBOX_TYPE": "langsmith"}),
            patch.dict(github_proxy.SANDBOX_BACKENDS, {"thread-1": backend}, clear=True),
            patch(
                "agent.utils.github_proxy.get_github_app_installation_token_with_expiry",
                new=AsyncMock(return_value=("ghs_new", new_expiry)),
            ),
            patch("agent.integrations.langsmith._configure_github_proxy") as mock_configure,
        ):
            result = await maybe_refresh_proxy_token("thread-1", now=now)

        assert result is True
        mock_configure.assert_called_once_with("sb-1", "ghs_new")
        expires_at, _recorded, _scope = github_proxy._PROXY_TOKEN_EXPIRY["thread-1"]
        assert expires_at == datetime(2025, 1, 1, 13, 0, 0, tzinfo=UTC)

    @pytest.mark.asyncio
    async def test_preserves_repo_scope_on_refresh(self) -> None:
        now = datetime(2025, 1, 1, 12, 0, 0, tzinfo=UTC)
        record_proxy_token_expiry("thread-1", now + timedelta(minutes=1), repositories=["open-swe"])
        backend = MagicMock(id="sb-1")
        token_mock = AsyncMock(return_value=("ghs_new", "2025-01-01T13:00:00Z"))

        with (
            patch.dict("os.environ", {"SANDBOX_TYPE": "langsmith"}),
            patch.dict(github_proxy.SANDBOX_BACKENDS, {"thread-1": backend}, clear=True),
            patch(
                "agent.utils.github_proxy.get_github_app_installation_token_with_expiry",
                new=token_mock,
            ),
            patch("agent.integrations.langsmith._configure_github_proxy"),
        ):
            result = await maybe_refresh_proxy_token("thread-1", now=now)

        assert result is True
        token_mock.assert_awaited_once_with(repositories=["open-swe"])
        _expires, _recorded, scope = github_proxy._PROXY_TOKEN_EXPIRY["thread-1"]
        assert scope == ("open-swe",)

    @pytest.mark.asyncio
    async def test_no_refresh_when_token_unavailable(self) -> None:
        now = datetime(2025, 1, 1, 12, 0, 0, tzinfo=UTC)
        record_proxy_token_expiry("thread-1", now + timedelta(minutes=1))
        backend = MagicMock(id="sb-1")

        with (
            patch.dict("os.environ", {"SANDBOX_TYPE": "langsmith"}),
            patch.dict(github_proxy.SANDBOX_BACKENDS, {"thread-1": backend}, clear=True),
            patch(
                "agent.utils.github_proxy.get_github_app_installation_token_with_expiry",
                new=AsyncMock(return_value=(None, None)),
            ),
            patch("agent.integrations.langsmith._configure_github_proxy") as mock_configure,
        ):
            result = await maybe_refresh_proxy_token("thread-1", now=now)

        assert result is False
        mock_configure.assert_not_called()


class TestRefreshGithubProxyMiddleware:
    @pytest.mark.asyncio
    async def test_calls_refresh_with_thread_id(self) -> None:
        from agent.middleware.refresh_github_proxy import refresh_github_proxy_before_model

        with (
            patch(
                "agent.middleware.refresh_github_proxy.get_config",
                return_value={"configurable": {"thread_id": "thread-9"}},
            ),
            patch(
                "agent.middleware.refresh_github_proxy.maybe_refresh_proxy_token",
                new=AsyncMock(return_value=True),
            ) as mock_refresh,
        ):
            result = await refresh_github_proxy_before_model.abefore_model({}, MagicMock())

        assert result is None
        mock_refresh.assert_awaited_once_with("thread-9")

    @pytest.mark.asyncio
    async def test_no_thread_id_is_noop(self) -> None:
        from agent.middleware.refresh_github_proxy import refresh_github_proxy_before_model

        with (
            patch(
                "agent.middleware.refresh_github_proxy.get_config",
                return_value={"configurable": {}},
            ),
            patch(
                "agent.middleware.refresh_github_proxy.maybe_refresh_proxy_token",
                new=AsyncMock(),
            ) as mock_refresh,
        ):
            result = await refresh_github_proxy_before_model.abefore_model({}, MagicMock())

        assert result is None
        mock_refresh.assert_not_called()

    @pytest.mark.asyncio
    async def test_swallows_refresh_errors(self) -> None:
        from agent.middleware.refresh_github_proxy import refresh_github_proxy_before_model

        with (
            patch(
                "agent.middleware.refresh_github_proxy.get_config",
                return_value={"configurable": {"thread_id": "thread-9"}},
            ),
            patch(
                "agent.middleware.refresh_github_proxy.maybe_refresh_proxy_token",
                new=AsyncMock(side_effect=RuntimeError("boom")),
            ),
        ):
            result = await refresh_github_proxy_before_model.abefore_model({}, MagicMock())

        assert result is None