ss-tools/backend/tests/test_api_key_auth.py

# #region TestAPIKeyAuth [C:3] [TYPE Module] [SEMANTICS test, api_key, auth, dependency]
# @BRIEF Contract tests for API key authentication — valid key, invalid, revoked, expired, missing permission,
#   environment scoping, JWT precedence. Uses TestClient with dependency overrides.
# @RELATION BINDS_TO -> [APIKeyUtilities]
# @RELATION BINDS_TO -> [APIKeyModel]
# @RELATION BINDS_TO -> [MaintenanceRoutesModule]
# @TEST_CONTRACT: get_api_key_principal returns APIKeyPrincipal for valid key, None for no header, 401 for invalid/revoked/expired.
# @TEST_CONTRACT: require_api_key_or_jwt rejects invalid/revoked/expired keys, checks permissions, enforces environment scope.
# @TEST_EDGE: missing_header -> returns None
# @TEST_EDGE: invalid_key -> 401
# @TEST_EDGE: revoked_key -> 401
# @TEST_EDGE: expired_key -> 401
# @TEST_EDGE: missing_permission -> 403
# @TEST_EDGE: environment_scope_mismatch -> 400
# @TEST_EDGE: jwt_precedence -> JWT takes precedence over API key
from datetime import UTC, datetime, timedelta
import pytest
from unittest.mock import AsyncMock, MagicMock

from fastapi import HTTPException
from fastapi.testclient import TestClient

# ── Fixtures ──────────────────────────────────────────────────


@pytest.fixture
def db_session():
    """Create an in-memory SQLite session with the api_keys table."""
    from sqlalchemy import create_engine
    from sqlalchemy.orm import sessionmaker
    from sqlalchemy.pool import StaticPool

    from src.models.mapping import Base

    engine = create_engine(
        "sqlite:///:memory:",
        poolclass=StaticPool,
        connect_args={"check_same_thread": False},
    )
    Base.metadata.create_all(engine)
    Session = sessionmaker(bind=engine)
    session = Session()
    yield session
    session.close()


@pytest.fixture
def valid_api_key(db_session):
    """Create a valid API key and return its raw value, prefix, and database row."""
    from src.core.auth.api_key import generate_api_key
    from src.models.api_key import APIKey

    raw, prefix, key_hash = generate_api_key()
    api_key = APIKey(
        key_hash=key_hash,
        prefix=prefix,
        name="Test Key",
        permissions=["maintenance:start", "maintenance:end"],
        active=True,
    )
    db_session.add(api_key)
    db_session.commit()
    return raw, api_key


@pytest.fixture
def client():
    """Create a TestClient."""
    from src.app import app

    return TestClient(app)


@pytest.fixture
def mock_db():
    """Patch get_db dependency with in-memory SQLite session (same pattern as test_maintenance_api)."""
    from sqlalchemy import create_engine
    from sqlalchemy.orm import sessionmaker
    from sqlalchemy.pool import StaticPool

    from src.app import app
    from src.dependencies import get_db
    from src.models.mapping import Base

    engine = create_engine(
        "sqlite:///:memory:",
        poolclass=StaticPool,
        connect_args={"check_same_thread": False},
    )
    Base.metadata.create_all(engine)
    Session = sessionmaker(bind=engine)
    session = Session()

    # Add default maintenance settings (needed by /api/maintenance endpoints)
    from src.models.maintenance import DashboardScope, MaintenanceSettings

    settings = MaintenanceSettings(
        id="default",
        target_environment_id="test-env",
        display_timezone="UTC",
        banner_template="Test: {message} ({start_time}-{end_time})",
        dashboard_scope=DashboardScope.PUBLISHED_ONLY,
        excluded_dashboard_ids=[],
        forced_dashboard_ids=[],
    )
    session.add(settings)
    session.commit()

    def _get_db_override():
        db = Session()
        try:
            yield db
        finally:
            db.close()

    app.dependency_overrides[get_db] = _get_db_override
    yield session
    app.dependency_overrides.pop(get_db, None)
    session.close()


@pytest.fixture
def mock_task_manager():
    """Patch get_task_manager to return a mock."""
    from src.app import app
    from src.dependencies import get_task_manager

    mock_tm = MagicMock()
    mock_task = MagicMock()
    mock_task.id = "test-task-id-123"
    mock_tm.create_task = AsyncMock(return_value=mock_task)

    app.dependency_overrides[get_task_manager] = lambda: mock_tm
    yield mock_tm
    app.dependency_overrides.pop(get_task_manager, None)


# ── Tests for get_api_key_principal ───────────────────────────

class TestGetApiKeyPrincipal:
    """Contract tests for get_api_key_principal dependency."""

    # #region test_no_header_returns_none [C:2] [TYPE Function]
    # @BRIEF No X-API-Key header → returns None.
    @pytest.mark.asyncio
    async def test_no_header_returns_none(self):
        from fastapi import Request

        from src.dependencies import get_api_key_principal

        # Mock a request without the header
        mock_request = MagicMock(spec=Request)
        mock_request.headers.get.return_value = None
        mock_db = MagicMock()

        result = await get_api_key_principal(mock_request, mock_db)
        assert result is None

    # #endregion test_no_header_returns_none

    # #region test_valid_key_returns_principal [C:2] [TYPE Function]
    # @BRIEF Valid API key → returns APIKeyPrincipal with correct fields.
    @pytest.mark.asyncio
    async def test_valid_key_returns_principal(self, db_session, valid_api_key):
        from fastapi import Request

        from src.dependencies import get_api_key_principal

        raw_key, api_key_row = valid_api_key
        mock_request = MagicMock(spec=Request)
        mock_request.headers.get.return_value = raw_key

        result = await get_api_key_principal(mock_request, db_session)
        assert result is not None
        assert result.name == "Test Key"
        assert result.api_key_id == api_key_row.id
        assert result.environment_id is None
        assert "maintenance:start" in result.permissions

    # #endregion test_valid_key_returns_principal

    # #region test_invalid_key_raises_401 [C:2] [TYPE Function]
    # @BRIEF Invalid API key → 401.
    @pytest.mark.asyncio
    async def test_invalid_key_raises_401(self, db_session):
        from fastapi import Request

        from src.dependencies import get_api_key_principal

        mock_request = MagicMock(spec=Request)
        mock_request.headers.get.return_value = "ssk_invalid_key_that_does_not_exist"

        with pytest.raises(HTTPException) as exc:
            await get_api_key_principal(mock_request, db_session)
        assert exc.value.status_code == 401

    # #endregion test_invalid_key_raises_401

    # #region test_revoked_key_raises_401 [C:2] [TYPE Function]
    # @BRIEF Revoked (active=False) API key → 401.
    @pytest.mark.asyncio
    async def test_revoked_key_raises_401(self, db_session):
        from fastapi import Request

        from src.core.auth.api_key import generate_api_key
        from src.dependencies import get_api_key_principal
        from src.models.api_key import APIKey

        raw, prefix, key_hash = generate_api_key()
        api_key = APIKey(
            key_hash=key_hash,
            prefix=prefix,
            name="Revoked Key",
            permissions=["maintenance:start"],
            active=False,  # revoked
        )
        db_session.add(api_key)
        db_session.commit()

        mock_request = MagicMock(spec=Request)
        mock_request.headers.get.return_value = raw

        with pytest.raises(HTTPException) as exc:
            await get_api_key_principal(mock_request, db_session)
        assert exc.value.status_code == 401

    # #endregion test_revoked_key_raises_401

    # #region test_expired_key_raises_401 [C:2] [TYPE Function]
    # @BRIEF Expired API key → 401.
    @pytest.mark.asyncio
    async def test_expired_key_raises_401(self, db_session):
        from fastapi import Request

        from src.core.auth.api_key import generate_api_key
        from src.dependencies import get_api_key_principal
        from src.models.api_key import APIKey

        raw, prefix, key_hash = generate_api_key()
        api_key = APIKey(
            key_hash=key_hash,
            prefix=prefix,
            name="Expired Key",
            permissions=["maintenance:start"],
            active=True,
            expires_at=datetime.now(UTC) - timedelta(hours=1),  # expired
        )
        db_session.add(api_key)
        db_session.commit()

        mock_request = MagicMock(spec=Request)
        mock_request.headers.get.return_value = raw

        with pytest.raises(HTTPException) as exc:
            await get_api_key_principal(mock_request, db_session)
        assert exc.value.status_code == 401

    # #endregion test_expired_key_raises_401


# #endregion TestAPIKeyAuth


# ── Tests for require_api_key_or_jwt (integration via TestClient) ──────

class TestRequireApiKeyOrJwt:
    """Contract tests for require_api_key_or_jwt dependency factory via TestClient."""

    API_START_URL = "/api/maintenance/start"
    START_PAYLOAD = {
        "tables": ["raw.sales"],
        "start_time": (datetime.now(UTC) + timedelta(hours=1)).isoformat(),
        "end_time": (datetime.now(UTC) + timedelta(hours=3)).isoformat(),
        "message": "Test maintenance",
        "environment_id": "ss-dev",
    }

    @staticmethod
    def _create_api_key(db, environment_id=None, permissions=None):
        """Helper to create an API key in the test DB and return the raw key."""
        from src.core.auth.api_key import generate_api_key
        from src.models.api_key import APIKey

        raw, prefix, key_hash = generate_api_key()
        api_key = APIKey(
            key_hash=key_hash,
            prefix=prefix,
            name="Test Key",
            permissions=permissions or ["maintenance:start"],
            active=True,
            environment_id=environment_id,
        )
        db.add(api_key)
        db.commit()
        return raw

    # #region test_api_key_auth_valid [C:2] [TYPE Function]
    # @BRIEF Valid API key with matching permission → 202.
    def test_api_key_auth_valid(self, client, mock_db, mock_task_manager):
        raw_key = self._create_api_key(mock_db, permissions=["maintenance:start"])

        response = client.post(
            self.API_START_URL,
            json=self.START_PAYLOAD,
            headers={"X-API-Key": raw_key},
        )
        assert response.status_code == 202
        data = response.json()
        assert "task_id" in data
        assert data["status"] == "pending"

    # #endregion test_api_key_auth_valid

    # #region test_api_key_auth_wrong_permission [C:2] [TYPE Function]
    # @BRIEF API key lacks required permission → 403.
    def test_api_key_auth_wrong_permission(self, client, mock_db, mock_task_manager):
        # Key has 'maintenance:end' only, not 'maintenance:start'
        raw_key = self._create_api_key(mock_db, permissions=["maintenance:end"])

        response = client.post(
            self.API_START_URL,
            json=self.START_PAYLOAD,
            headers={"X-API-Key": raw_key},
        )
        assert response.status_code == 403
        data = response.json()
        assert "detail" in data

    # #endregion test_api_key_auth_wrong_permission

    # #region test_api_key_auth_revoked [C:2] [TYPE Function]
    # @BRIEF Revoked API key → 401.
    def test_api_key_auth_revoked(self, client, mock_db, mock_task_manager):
        from src.core.auth.api_key import generate_api_key
        from src.models.api_key import APIKey

        raw, prefix, key_hash = generate_api_key()
        api_key = APIKey(
            key_hash=key_hash,
            prefix=prefix,
            name="Revoked Key",
            permissions=["maintenance:start"],
            active=False,
        )
        mock_db.add(api_key)
        mock_db.commit()

        response = client.post(
            self.API_START_URL,
            json=self.START_PAYLOAD,
            headers={"X-API-Key": raw},
        )
        assert response.status_code == 401
        assert "revoked" in response.json().get("detail", "").lower()

    # #endregion test_api_key_auth_revoked

    # #region test_api_key_auth_expired [C:2] [TYPE Function]
    # @BRIEF Expired API key → 401.
    def test_api_key_auth_expired(self, client, mock_db, mock_task_manager):
        from src.core.auth.api_key import generate_api_key
        from src.models.api_key import APIKey

        raw, prefix, key_hash = generate_api_key()
        api_key = APIKey(
            key_hash=key_hash,
            prefix=prefix,
            name="Expired Key",
            permissions=["maintenance:start"],
            active=True,
            expires_at=datetime.now(UTC) - timedelta(hours=1),
        )
        mock_db.add(api_key)
        mock_db.commit()

        response = client.post(
            self.API_START_URL,
            json=self.START_PAYLOAD,
            headers={"X-API-Key": raw},
        )
        assert response.status_code == 401
        assert "expired" in response.json().get("detail", "").lower()

    # #endregion test_api_key_auth_expired

    # #region test_api_key_auth_invalid [C:2] [TYPE Function]
    # @BRIEF Non-existent key → 401.
    def test_api_key_auth_invalid(self, client, mock_db, mock_task_manager):
        response = client.post(
            self.API_START_URL,
            json=self.START_PAYLOAD,
            headers={"X-API-Key": "ssk_invalid_nonexistent_key"},
        )
        assert response.status_code == 401
        assert "invalid" in response.json().get("detail", "").lower()

    # #endregion test_api_key_auth_invalid

    # #region test_api_key_auth_environment_scope [C:2] [TYPE Function]
    # @BRIEF Key scoped to ss-dev, request for ss-prod → 400.
    def test_api_key_auth_environment_scope(self, client, mock_db, mock_task_manager):
        raw_key = self._create_api_key(
            mock_db,
            environment_id="ss-dev",
            permissions=["maintenance:start"],
        )

        payload = dict(self.START_PAYLOAD)
        payload["environment_id"] = "ss-prod"  # Mismatch!

        response = client.post(
            self.API_START_URL,
            json=payload,
            headers={"X-API-Key": raw_key},
        )
        assert response.status_code == 400
        data = response.json()
        assert "restricted" in data.get("detail", "").lower()
        assert "ss-dev" in data.get("detail", "")

    # #endregion test_api_key_auth_environment_scope

    # #region test_api_key_auth_environment_scope_ok [C:2] [TYPE Function]
    # @BRIEF Key scoped to ss-dev, request for ss-dev → 202.
    def test_api_key_auth_environment_scope_ok(self, client, mock_db, mock_task_manager):
        raw_key = self._create_api_key(
            mock_db,
            environment_id="ss-dev",
            permissions=["maintenance:start"],
        )

        payload = dict(self.START_PAYLOAD)
        payload["environment_id"] = "ss-dev"  # Match!

        response = client.post(
            self.API_START_URL,
            json=payload,
            headers={"X-API-Key": raw_key},
        )
        assert response.status_code == 202
        data = response.json()
        assert "task_id" in data

    # #endregion test_api_key_auth_environment_scope_ok

    # #region test_jwt_precedence [C:2] [TYPE Function]
    # @BRIEF Both valid API key + valid JWT → JWT takes precedence.
    def test_jwt_precedence(self, client, mock_db, mock_task_manager):
        from src.app import app
        from src.core.auth.jwt import create_access_token
        from src.dependencies import get_auth_db
        from src.models.auth import Role, User

        raw_key = self._create_api_key(
            mock_db,
            environment_id="ss-dev",
            permissions=["maintenance:start"],
        )

        # Create a JWT user with Admin role (full access)
        role = Role(name="Admin", description="Admin role")
        mock_db.add(role)
        mock_db.commit()

        user = User(
            username="jwt-user",
            password_hash="nohash",
            auth_source="LOCAL",
            is_active=True,
        )
        user.roles.append(role)
        mock_db.add(user)
        mock_db.commit()

        # Create a valid JWT token
        token = create_access_token(data={"sub": "jwt-user"})

        # Override get_auth_db to use our mock_db for auth queries too
        # The auth DB and main DB are separate in production, but for tests
        # we use the same in-memory DB

        def _get_auth_db_override():
            yield mock_db

        app.dependency_overrides[get_auth_db] = _get_auth_db_override

        # Send request with BOTH X-API-Key and Authorization: Bearer
        response = client.post(
            self.API_START_URL,
            json=self.START_PAYLOAD,
            headers={
                "X-API-Key": raw_key,
                "Authorization": f"Bearer {token}",
            },
        )
        assert response.status_code == 202
        data = response.json()
        assert "task_id" in data

        # Cleanup
        app.dependency_overrides.pop(get_auth_db, None)

    # #endregion test_jwt_precedence


# #endregion TestRequireApiKeyOrJwt