open-swe/tests/test_team_credentials.py

from __future__ import annotations

from typing import Any

import pytest
from cryptography.fernet import Fernet
from pydantic import ValidationError

from agent.dashboard import team_credentials as tc
from agent.dashboard.team_credentials import (
    DatadogCredentialsUpdate,
    LangSmithCredentialsUpdate,
)


class _FakeStore:
    def __init__(self) -> None:
        self.items: dict[tuple[tuple[str, ...], str], dict[str, Any]] = {}

    async def get_item(self, namespace: list[str], key: str):
        value = self.items.get((tuple(namespace), key))
        return {"value": value} if value is not None else None

    async def put_item(self, namespace: list[str], key: str, value: dict[str, Any]) -> None:
        self.items[(tuple(namespace), key)] = value

    async def delete_item(self, namespace: list[str], key: str) -> None:
        self.items.pop((tuple(namespace), key), None)


class _FakeClient:
    def __init__(self, store: _FakeStore) -> None:
        self.store = store


@pytest.fixture()
def fake_store(monkeypatch: pytest.MonkeyPatch) -> _FakeStore:
    store = _FakeStore()
    monkeypatch.setattr(tc, "_client", lambda: _FakeClient(store))
    monkeypatch.setenv("TOKEN_ENCRYPTION_KEY", Fernet.generate_key().decode())
    return store


class TestValidators:
    def test_site_normalized(self) -> None:
        u = DatadogCredentialsUpdate(site="https://app.datadoghq.com/", api_key="k", app_key="a")
        assert u.site == "datadoghq.com"

    def test_site_default(self) -> None:
        u = DatadogCredentialsUpdate(api_key="k", app_key="a")
        assert u.site == "datadoghq.com"

    def test_unsupported_site_rejected(self) -> None:
        with pytest.raises(ValidationError):
            DatadogCredentialsUpdate(site="evil.example.com", api_key="k", app_key="a")

    def test_empty_keys_rejected(self) -> None:
        with pytest.raises(ValidationError):
            DatadogCredentialsUpdate(api_key="  ", app_key="a")
        with pytest.raises(ValidationError):
            LangSmithCredentialsUpdate(api_key="")

    def test_langsmith_endpoint_normalized(self) -> None:
        u = LangSmithCredentialsUpdate(api_key="k", endpoint="https://x/")
        assert u.endpoint == "https://x"


class TestMcpUrl:
    def test_url_includes_site_and_toolsets(self) -> None:
        creds = tc.DatadogCredentials(site="datadoghq.eu", api_key="a", app_key="b")
        assert creds.mcp_url("core") == (
            "https://mcp.datadoghq.eu/api/unstable/mcp-server/mcp?toolsets=core"
        )


@pytest.mark.asyncio
async def test_datadog_roundtrip_and_redaction(fake_store: _FakeStore) -> None:
    status = await tc.connect_datadog(
        DatadogCredentialsUpdate(site="datadoghq.com", api_key="secret-api-1234", app_key="appk")
    )
    assert status["datadog"]["connected"] is True
    assert status["datadog"]["api_key_last4"] == "1234"
    # Stored record holds ciphertext, not the plaintext key.
    record = fake_store.items[(("team_credentials",), "datadog")]
    assert record["encrypted_api_key"] != "secret-api-1234"

    creds = await tc.get_datadog_credentials()
    assert creds is not None
    assert creds.api_key == "secret-api-1234"
    assert creds.app_key == "appk"

    after = await tc.disconnect_datadog()
    assert after["datadog"]["connected"] is False
    assert await tc.get_datadog_credentials() is None


@pytest.mark.asyncio
async def test_langsmith_roundtrip(fake_store: _FakeStore) -> None:
    status = await tc.connect_langsmith(LangSmithCredentialsUpdate(api_key="ls-key-9999"))
    assert status["langsmith"]["connected"] is True
    assert status["langsmith"]["api_key_last4"] == "9999"

    creds = await tc.get_langsmith_credentials()
    assert creds is not None
    assert creds.api_key == "ls-key-9999"
    assert creds.endpoint == tc.DEFAULT_LANGSMITH_ENDPOINT

    await tc.disconnect_langsmith()
    assert await tc.get_langsmith_credentials() is None


@pytest.mark.asyncio
async def test_status_empty_when_unset(fake_store: _FakeStore) -> None:
    status = await tc.get_team_credentials_status()
    assert status["datadog"]["connected"] is False
    assert status["langsmith"]["connected"] is False


@pytest.mark.asyncio
async def test_connecting_one_keeps_other(fake_store: _FakeStore) -> None:
    await tc.connect_datadog(DatadogCredentialsUpdate(api_key="dd", app_key="app"))
    await tc.connect_langsmith(LangSmithCredentialsUpdate(api_key="ls"))
    status = await tc.get_team_credentials_status()
    assert status["datadog"]["connected"] is True
    assert status["langsmith"]["connected"] is True