fastmcp/tests/cli/deploy/test_command.py

import json
from collections.abc import Callable, Iterator
from contextlib import contextmanager
from pathlib import Path
from unittest.mock import Mock
from urllib.parse import parse_qs

import httpx2
import pytest
from pydantic import SecretStr

import fastmcp
import fastmcp.cli.deploy.authentication as authentication_module
import fastmcp.cli.deploy.command as command_module
from fastmcp.cli.deploy.command import login, logout, whoami
from fastmcp.cli.deploy.configuration import ConfigurationStore
from fastmcp.cli.deploy.credentials import CredentialStore
from fastmcp.cli.deploy.horizon_client import HorizonClient
from fastmcp.cli.deploy.state import StateFileError


class HorizonAuthAPI:
    def __init__(
        self,
        *,
        token_error: str | None = None,
        revoke_status: int = 204,
        invalid_api_key: str | None = None,
        on_request: Callable[[httpx2.Request], None] | None = None,
    ) -> None:
        self.token_error = token_error
        self.revoke_status = revoke_status
        self.invalid_api_key = invalid_api_key
        self.on_request = on_request
        self.requests: list[httpx2.Request] = []

    def __call__(self, request: httpx2.Request) -> httpx2.Response:
        self.requests.append(request)
        if self.on_request is not None:
            self.on_request(request)
        path = request.url.path
        if path == "/api/v0/oauth/device/authorization":
            return httpx2.Response(
                200,
                json={
                    "device_code": "device-secret",
                    "user_code": "ABCD-EFGH",
                    "verification_uri": "https://horizon.prefect.io/oauth/device",
                    "verification_uri_complete": (
                        "https://horizon.prefect.io/oauth/device?user_code=ABCD-EFGH"
                    ),
                    "expires_in": 600,
                    "interval": 1,
                },
            )
        if path == "/api/v0/oauth/device/token":
            if self.token_error is not None:
                return httpx2.Response(400, json={"error": self.token_error})
            return httpx2.Response(
                200,
                json={"access_token": "fmcp_device_key", "token_type": "Bearer"},
            )
        if path == "/api/v0/me":
            if request.headers.get("Authorization") == (
                f"Bearer {self.invalid_api_key}"
            ):
                return httpx2.Response(401)
            return httpx2.Response(
                200,
                json={
                    "user": {
                        "id": "user-1",
                        "email": "ada@example.com",
                        "name": "Ada",
                    }
                },
            )
        if path == "/api/v0/me/api-key":
            return httpx2.Response(self.revoke_status)
        raise AssertionError(f"Unexpected request: {request.method} {path}")


@pytest.fixture
def use_horizon_api(
    monkeypatch: pytest.MonkeyPatch,
) -> Callable[[HorizonAuthAPI], None]:
    def use(api: HorizonAuthAPI) -> None:
        transport = httpx2.MockTransport(api)

        def client(
            api_origin: str,
            *,
            api_key: SecretStr | str | None = None,
        ) -> HorizonClient:
            return HorizonClient(
                api_origin,
                api_key=api_key,
                transport=transport,
            )

        monkeypatch.setattr(command_module, "HorizonClient", client)

    return use


def test_session_snapshot_reads_host_and_credential_under_one_lock(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    events: list[str] = []
    configuration_load = ConfigurationStore.load
    credential_load = CredentialStore.load

    @contextmanager
    def lock(directory: Path) -> Iterator[None]:
        events.append("lock")
        yield
        events.append("unlock")

    def load_configuration(store: ConfigurationStore):
        events.append("configuration")
        return configuration_load(store)

    def load_credential(store: CredentialStore):
        events.append("credential")
        return credential_load(store)

    monkeypatch.setattr(command_module, "state_lock", lock)
    monkeypatch.setattr(ConfigurationStore, "load", load_configuration)
    monkeypatch.setattr(CredentialStore, "load", load_credential)

    configuration, credential = command_module._load_session_snapshot(CredentialStore())

    assert configuration.api_origin == "https://horizon.prefect.io"
    assert credential is None
    assert events == ["lock", "configuration", "credential", "unlock"]


@pytest.fixture(autouse=True)
def no_device_poll_delay(monkeypatch: pytest.MonkeyPatch) -> None:
    async def sleep(_: float) -> None:
        return None

    monkeypatch.setattr(authentication_module.asyncio, "sleep", sleep)


async def test_json_login_writes_one_result_and_challenge_to_stderr(
    use_horizon_api: Callable[[HorizonAuthAPI], None],
    capsys: pytest.CaptureFixture[str],
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    api = HorizonAuthAPI()
    use_horizon_api(api)
    browser_open = Mock()
    monkeypatch.setattr(command_module.webbrowser, "open", browser_open)
    monkeypatch.setattr(command_module.platform, "node", lambda: "Avery's laptop")
    monkeypatch.setattr(command_module.platform, "system", lambda: "Darwin")
    monkeypatch.setattr(command_module.platform, "machine", lambda: "arm64")
    monkeypatch.setattr(command_module.fastmcp, "__version__", "4.0.0")

    await login(json_output=True)

    captured = capsys.readouterr()
    stdout_lines = captured.out.strip().splitlines()
    assert len(stdout_lines) == 1
    assert json.loads(stdout_lines[0]) == {
        "ok": True,
        "command": "login",
        "user": {
            "id": "user-1",
            "email": "ada@example.com",
            "name": "Ada",
        },
    }
    assert json.loads(captured.err) == {
        "event": "device_authorization",
        "verificationUrl": "https://horizon.prefect.io/oauth/device",
        "verificationUrlComplete": (
            "https://horizon.prefect.io/oauth/device?user_code=ABCD-EFGH"
        ),
        "userCode": "ABCD-EFGH",
    }
    browser_open.assert_not_called()
    authorization_request = next(
        request
        for request in api.requests
        if request.url.path == "/api/v0/oauth/device/authorization"
    )
    assert parse_qs(authorization_request.content.decode()) == {
        "client_id": ["fastmcp-cli"],
        "device_name": ["Avery's laptop"],
        "platform": ["darwin"],
        "architecture": ["arm64"],
        "client_version": ["4.0.0"],
    }

    state = json.loads(CredentialStore().path.read_text())
    assert state == {"schemaVersion": 1, "apiKey": "fmcp_device_key"}
    assert not (fastmcp.settings.home / "cli" / "config.json").exists()


async def test_login_host_is_saved_before_device_authorization(
    use_horizon_api: Callable[[HorizonAuthAPI], None],
    capsys: pytest.CaptureFixture[str],
) -> None:
    api = HorizonAuthAPI()
    use_horizon_api(api)

    await login(host="https://dev.horizon.prefect.io/", json_output=True)

    assert json.loads(capsys.readouterr().out)["ok"] is True
    configuration_path = fastmcp.settings.home / "cli" / "config.json"
    assert json.loads(configuration_path.read_text()) == {
        "schemaVersion": 1,
        "apiOrigin": "https://dev.horizon.prefect.io",
    }
    assert {request.url.host for request in api.requests} == {"dev.horizon.prefect.io"}


async def test_login_rejects_an_invalid_host(
    capsys: pytest.CaptureFixture[str],
) -> None:
    with pytest.raises(SystemExit, match="1"):
        await login(host="https://horizon.prefect.io/path", json_output=True)

    result = json.loads(capsys.readouterr().out)
    assert result["error"]["category"] == "invalid_host"
    assert CredentialStore().path.exists() is False


async def test_tty_login_survives_browser_open_failure(
    use_horizon_api: Callable[[HorizonAuthAPI], None],
    capsys: pytest.CaptureFixture[str],
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    use_horizon_api(HorizonAuthAPI())
    browser_open = Mock(side_effect=OSError("No browser"))
    monkeypatch.setattr(command_module, "_can_open_browser", lambda: True)
    monkeypatch.setattr(command_module.webbrowser, "open", browser_open)

    await login()

    output = capsys.readouterr().out
    assert "https://horizon.prefect.io/oauth/device" in output
    assert "ABCD-EFGH" in output
    assert "Logged into Horizon" in output
    assert "Ada" in output
    assert "ada@example.com" in output
    assert "Organization" not in output
    browser_open.assert_called_once()


async def test_whoami_uses_the_stored_key_after_a_restart(
    use_horizon_api: Callable[[HorizonAuthAPI], None],
    capsys: pytest.CaptureFixture[str],
) -> None:
    api = HorizonAuthAPI()
    use_horizon_api(api)
    await login(json_output=True)
    capsys.readouterr()

    await whoami(json_output=True)

    result = json.loads(capsys.readouterr().out)
    assert result["command"] == "whoami"
    assert result["user"]["email"] == "ada@example.com"
    assert [request.url.path for request in api.requests].count("/api/v0/me") == 2


async def test_login_replaces_an_invalid_stored_key(
    use_horizon_api: Callable[[HorizonAuthAPI], None],
    capsys: pytest.CaptureFixture[str],
) -> None:
    use_horizon_api(HorizonAuthAPI(invalid_api_key="fmcp_stale_key"))
    CredentialStore().save("fmcp_stale_key")

    await login(json_output=True)

    captured = capsys.readouterr()
    assert json.loads(captured.out)["ok"] is True
    assert json.loads(captured.err)["event"] == "device_authorization"
    stored_key = CredentialStore().load()
    assert stored_key is not None
    assert stored_key.get_secret_value() == "fmcp_device_key"


async def test_login_never_persists_an_environment_key(
    use_horizon_api: Callable[[HorizonAuthAPI], None],
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    api = HorizonAuthAPI()
    use_horizon_api(api)
    monkeypatch.setenv("HORIZON_API_KEY", "fmcp_environment_key")

    await login(json_output=True)

    assert CredentialStore().path.exists() is False
    assert not any(
        request.url.path.startswith("/api/v0/oauth/device") for request in api.requests
    )


async def test_json_whoami_reports_a_failed_rejected_key_cleanup(
    use_horizon_api: Callable[[HorizonAuthAPI], None],
    capsys: pytest.CaptureFixture[str],
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    use_horizon_api(HorizonAuthAPI(invalid_api_key="fmcp_stale_key"))
    CredentialStore().save("fmcp_stale_key")

    def fail_clear(store: CredentialStore) -> None:
        raise StateFileError("cleanup failed")

    monkeypatch.setattr(CredentialStore, "clear", fail_clear)

    with pytest.raises(SystemExit, match="1"):
        await whoami(json_output=True)

    result = json.loads(capsys.readouterr().out)
    assert result["error"]["category"] == "state_error"


async def test_whoami_does_not_clear_a_newer_host_credential(
    use_horizon_api: Callable[[HorizonAuthAPI], None],
    capsys: pytest.CaptureFixture[str],
) -> None:
    credentials = CredentialStore()
    credentials.save("fmcp_stale_key")
    switched = False

    def switch_host(request: httpx2.Request) -> None:
        nonlocal switched
        if request.url.path != "/api/v0/me" or switched:
            return
        switched = True
        ConfigurationStore().set_api_origin(
            "https://dev.horizon.prefect.io",
            credentials=credentials,
        )
        credentials.save_for_origin(
            "fmcp_new_key",
            expected_api_origin="https://dev.horizon.prefect.io",
        )

    api = HorizonAuthAPI(
        invalid_api_key="fmcp_stale_key",
        on_request=switch_host,
    )
    use_horizon_api(api)

    with pytest.raises(SystemExit, match="1"):
        await whoami(json_output=True)

    result = json.loads(capsys.readouterr().out)
    assert result["error"]["category"] == "authentication_invalid"
    assert ConfigurationStore().load().api_origin == ("https://dev.horizon.prefect.io")
    stored_key = credentials.load()
    assert stored_key is not None
    assert stored_key.get_secret_value() == "fmcp_new_key"
    assert api.requests[0].url.host == "horizon.prefect.io"


async def test_json_whoami_does_not_start_device_authorization(
    capsys: pytest.CaptureFixture[str],
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    browser_open = Mock()
    monkeypatch.setattr(command_module.webbrowser, "open", browser_open)

    with pytest.raises(SystemExit, match="1"):
        await whoami(json_output=True)

    result = json.loads(capsys.readouterr().out)
    assert result["error"]["category"] == "authentication_required"
    browser_open.assert_not_called()


@pytest.mark.parametrize(
    ("token_error", "category"),
    [
        ("access_denied", "authorization_denied"),
        ("expired_token", "authorization_expired"),
    ],
)
async def test_json_login_reports_stable_device_failures(
    token_error: str,
    category: str,
    use_horizon_api: Callable[[HorizonAuthAPI], None],
    capsys: pytest.CaptureFixture[str],
) -> None:
    use_horizon_api(HorizonAuthAPI(token_error=token_error))

    with pytest.raises(SystemExit, match="1"):
        await login(json_output=True)

    result = json.loads(capsys.readouterr().out)
    assert result["error"]["category"] == category
    assert CredentialStore().path.exists() is False


async def test_logout_does_not_modify_environment_or_stored_credentials(
    use_horizon_api: Callable[[HorizonAuthAPI], None],
    capsys: pytest.CaptureFixture[str],
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    api = HorizonAuthAPI()
    use_horizon_api(api)
    monkeypatch.setenv("HORIZON_API_KEY", "fmcp_environment_key")
    CredentialStore().save("fmcp_stored_key")

    await logout(json_output=True)

    assert json.loads(capsys.readouterr().out) == {
        "ok": True,
        "command": "logout",
        "credentialSource": "environment",
        "localCredentialRemoved": False,
        "remoteRevoked": False,
    }
    stored_key = CredentialStore().load()
    assert stored_key is not None
    assert stored_key.get_secret_value() == "fmcp_stored_key"
    assert api.requests == []


async def test_logout_revokes_the_remote_key_and_clears_local_state(
    use_horizon_api: Callable[[HorizonAuthAPI], None],
    capsys: pytest.CaptureFixture[str],
) -> None:
    api = HorizonAuthAPI()
    use_horizon_api(api)
    CredentialStore().save("fmcp_stored_key")

    await logout(json_output=True)

    assert json.loads(capsys.readouterr().out) == {
        "ok": True,
        "command": "logout",
        "localCredentialRemoved": True,
        "remoteRevoked": True,
    }
    assert CredentialStore().path.exists() is False
    assert any(
        request.method == "DELETE" and request.url.path == "/api/v0/me/api-key"
        for request in api.requests
    )


async def test_logout_clears_local_state_when_remote_revocation_fails(
    use_horizon_api: Callable[[HorizonAuthAPI], None],
    capsys: pytest.CaptureFixture[str],
) -> None:
    use_horizon_api(HorizonAuthAPI(revoke_status=503))
    CredentialStore().save("fmcp_stored_key")

    with pytest.raises(SystemExit, match="1"):
        await logout(json_output=True)

    result = json.loads(capsys.readouterr().out)
    assert result["error"]["category"] == "remote_revocation_failed"
    assert result["localCredentialRemoved"] is True
    assert result["remoteCredentialMayRemain"] is True
    assert CredentialStore().path.exists() is False