konstruct/tests/unit/test_llm_keys_crud.py

"""
Unit tests for LLM key CRUD API endpoints.

Tests:
  - test_create_llm_key: POST encrypts key and returns {id, provider, label, key_hint, created_at}
    with no api_key in the response
  - test_list_llm_keys_redacted: GET returns list without full key, only key_hint (last 4 chars)
  - test_delete_llm_key: DELETE removes key, subsequent GET no longer includes it
  - test_create_duplicate_provider: POST with same tenant_id+provider returns 409 Conflict
  - test_delete_nonexistent_key: DELETE with unknown key_id returns 404
"""

from __future__ import annotations

import uuid
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
from fastapi import HTTPException

from shared.api.llm_keys import (
    LlmKeyCreate,
    LlmKeyResponse,
    create_llm_key,
    delete_llm_key,
    list_llm_keys,
)
from shared.models.billing import TenantLlmKey


def _make_mock_key(
    tenant_id: uuid.UUID,
    provider: str,
    label: str,
    api_key: str,
) -> TenantLlmKey:
    """Create a mock TenantLlmKey instance."""
    key = MagicMock(spec=TenantLlmKey)
    key.id = uuid.uuid4()
    key.tenant_id = tenant_id
    key.provider = provider
    key.label = label
    key.encrypted_key = f"ENCRYPTED_{api_key}"
    key.key_hint = api_key[-4:]
    key.key_version = 1
    from datetime import datetime, timezone
    key.created_at = datetime.now(tz=timezone.utc)
    return key


@pytest.fixture()
def tenant_id() -> uuid.UUID:
    return uuid.uuid4()


@pytest.fixture()
def mock_session() -> AsyncMock:
    session = AsyncMock()
    session.execute.return_value = MagicMock(
        scalar_one_or_none=MagicMock(return_value=None),
        scalars=MagicMock(return_value=MagicMock(all=MagicMock(return_value=[]))),
    )
    return session


@pytest.mark.asyncio
async def test_create_llm_key(tenant_id: uuid.UUID, mock_session: AsyncMock) -> None:
    """POST encrypts the key and returns {id, provider, label, key_hint, created_at} — no api_key."""
    api_key = "sk-test-openai-key-ABCD"
    body = LlmKeyCreate(provider="openai", label="Production Key", api_key=api_key)

    # Mock: no existing key for this tenant+provider
    mock_session.execute.return_value = MagicMock(
        scalar_one_or_none=MagicMock(return_value=None)
    )

    created_key = _make_mock_key(tenant_id, "openai", "Production Key", api_key)

    # Intercept session.add to capture the key and inject the mock
    def _capture_add(obj: object) -> None:
        # Mimic the ORM populating the key attributes after add+refresh
        pass

    mock_session.add.side_effect = _capture_add

    # Patch refresh to update the "obj" with values from created_key
    async def _mock_refresh(obj: object) -> None:
        obj.id = created_key.id  # type: ignore[union-attr]
        obj.created_at = created_key.created_at  # type: ignore[union-attr]

    mock_session.refresh = AsyncMock(side_effect=_mock_refresh)

    with patch("shared.api.llm_keys._get_encryption_service") as mock_enc:
        enc = MagicMock()
        enc.encrypt.return_value = "FERNET_CIPHERTEXT"
        mock_enc.return_value = enc

        with patch("shared.api.llm_keys._get_tenant_or_404", new_callable=AsyncMock):
            response = await create_llm_key(
                tenant_id=tenant_id,
                body=body,
                session=mock_session,
            )

    assert isinstance(response, LlmKeyResponse)
    assert response.provider == "openai"
    assert response.label == "Production Key"
    assert response.key_hint == api_key[-4:]
    # Response must NOT contain the raw api_key
    response_dict = response.model_dump()
    assert "api_key" not in response_dict
    assert "encrypted_key" not in response_dict


@pytest.mark.asyncio
async def test_list_llm_keys_redacted(tenant_id: uuid.UUID, mock_session: AsyncMock) -> None:
    """GET returns list with key_hint (last 4 chars) — never full key or encrypted_key."""
    api_key1 = "sk-openai-secret-WXYZ"
    api_key2 = "sk-anthropic-key-1234"

    key1 = _make_mock_key(tenant_id, "openai", "OpenAI Key", api_key1)
    key2 = _make_mock_key(tenant_id, "anthropic", "Anthropic Key", api_key2)

    mock_session.execute.return_value = MagicMock(
        scalars=MagicMock(return_value=MagicMock(all=MagicMock(return_value=[key1, key2])))
    )

    with patch("shared.api.llm_keys._get_tenant_or_404", new_callable=AsyncMock):
        response = await list_llm_keys(tenant_id=tenant_id, session=mock_session)

    assert len(response) == 2
    for item in response:
        assert isinstance(item, LlmKeyResponse)
        item_dict = item.model_dump()
        # Must not contain full key or encrypted data
        assert "api_key" not in item_dict
        assert "encrypted_key" not in item_dict
        # Must contain key_hint
        assert item_dict["key_hint"] is not None
        assert len(item_dict["key_hint"]) <= 4


@pytest.mark.asyncio
async def test_delete_llm_key(tenant_id: uuid.UUID, mock_session: AsyncMock) -> None:
    """DELETE removes the key and returns 204-equivalent (None)."""
    key_id = uuid.uuid4()
    existing_key = _make_mock_key(tenant_id, "openai", "Key to delete", "sk-DELETE-ABCD")

    mock_session.execute.return_value = MagicMock(
        scalar_one_or_none=MagicMock(return_value=existing_key)
    )

    with patch("shared.api.llm_keys._get_tenant_or_404", new_callable=AsyncMock):
        result = await delete_llm_key(
            tenant_id=tenant_id,
            key_id=key_id,
            session=mock_session,
        )

    assert result is None
    mock_session.delete.assert_called_once_with(existing_key)
    mock_session.commit.assert_called_once()


@pytest.mark.asyncio
async def test_create_duplicate_provider(tenant_id: uuid.UUID, mock_session: AsyncMock) -> None:
    """POST with same tenant_id+provider returns 409 Conflict."""
    # Mock: existing key found for this tenant+provider
    existing_key = _make_mock_key(tenant_id, "openai", "Existing Key", "sk-EXISTING-9999")
    mock_session.execute.return_value = MagicMock(
        scalar_one_or_none=MagicMock(return_value=existing_key)
    )

    body = LlmKeyCreate(provider="openai", label="Duplicate Key", api_key="sk-NEW-KEY-0000")

    with patch("shared.api.llm_keys._get_tenant_or_404", new_callable=AsyncMock):
        with pytest.raises(HTTPException) as exc_info:
            await create_llm_key(
                tenant_id=tenant_id,
                body=body,
                session=mock_session,
            )

    assert exc_info.value.status_code == 409


@pytest.mark.asyncio
async def test_delete_nonexistent_key(tenant_id: uuid.UUID, mock_session: AsyncMock) -> None:
    """DELETE with unknown key_id returns 404."""
    key_id = uuid.uuid4()

    # Mock: no key found
    mock_session.execute.return_value = MagicMock(
        scalar_one_or_none=MagicMock(return_value=None)
    )

    with patch("shared.api.llm_keys._get_tenant_or_404", new_callable=AsyncMock):
        with pytest.raises(HTTPException) as exc_info:
            await delete_llm_key(
                tenant_id=tenant_id,
                key_id=key_id,
                session=mock_session,
            )

    assert exc_info.value.status_code == 404