mnemosyne/mnemosyne/mcp_server/tests/test_teams_api.py

"""Tests for the ``/mcp_server/api/teams/`` REST control plane.

This is the Daedalus-facing surface described in §7 of
``docs/DAEDALUS_PALLAS_INTEGRATION_v1.md``. We do NOT exercise the
DRF Token / Session auth machinery here (that's covered by DRF
itself); instead we use :meth:`APIClient.force_authenticate` to focus
on the endpoints' own idempotence, ownership, and state-transition
rules.
"""

from __future__ import annotations

import uuid

import jwt as pyjwt
from django.contrib.auth import get_user_model
from django.test import TestCase
from django.urls import reverse
from rest_framework import status
from rest_framework.test import APIClient

from mcp_server.models import (
    MCPSigningKey,
    Team,
    TeamWorkspaceAssignment,
)

User = get_user_model()


def _seed_signing_key() -> MCPSigningKey:
    return MCPSigningKey.objects.create(
        kid=f"test-{uuid.uuid4().hex[:6]}",
        secret_hex="a" * 64,
        is_active=True,
    )


class _AuthenticatedAPITest(TestCase):
    """Shared ``APIClient`` authenticated as a regular user.

    The endpoints scope by ``request.user`` (every team has an
    ``owner``); ``self.user`` is the team owner and ``self.other_user``
    is used for cross-user access tests.
    """

    @classmethod
    def setUpTestData(cls):
        cls.user = User.objects.create_user(username="alice", password="pw")
        cls.other_user = User.objects.create_user(
            username="bob", password="pw"
        )

    def setUp(self):
        self.client = APIClient()
        self.client.force_authenticate(user=self.user)


# ---------------------------------------------------------------------------
# POST /mcp_server/api/teams/
# ---------------------------------------------------------------------------


class TeamCreateTest(_AuthenticatedAPITest):
    def setUp(self):
        super().setUp()
        self.url = reverse("mcp-server-api:team-create")
        _seed_signing_key()

    def test_requires_authentication(self):
        self.client.force_authenticate(user=None)
        resp = self.client.post(
            self.url, {"id": str(uuid.uuid4()), "name": "t"}, format="json"
        )
        self.assertIn(resp.status_code, (401, 403))

    def test_rejects_missing_fields(self):
        resp = self.client.post(self.url, {}, format="json")
        self.assertEqual(resp.status_code, status.HTTP_400_BAD_REQUEST)
        self.assertIn("id", resp.data)
        self.assertIn("name", resp.data)

    def test_rejects_non_uuid_id(self):
        resp = self.client.post(
            self.url, {"id": "not-a-uuid", "name": "t"}, format="json"
        )
        self.assertEqual(resp.status_code, status.HTTP_400_BAD_REQUEST)

    def test_creates_team_and_returns_jwt(self):
        tid = uuid.uuid4()
        resp = self.client.post(
            self.url, {"id": str(tid), "name": "Harper"}, format="json"
        )
        self.assertEqual(resp.status_code, status.HTTP_201_CREATED)
        self.assertEqual(uuid.UUID(resp.data["id"]), tid)
        self.assertEqual(resp.data["name"], "Harper")
        self.assertTrue(resp.data["active"])
        self.assertIn("jwt", resp.data)

        # JWT decodes and carries the right sub + jti.
        team = Team.objects.get(pk=tid)
        header = pyjwt.get_unverified_header(resp.data["jwt"])
        key = MCPSigningKey.objects.by_kid(header["kid"])
        decoded = pyjwt.decode(
            resp.data["jwt"],
            bytes.fromhex(key.secret_hex),
            algorithms=["HS256"],
            options={"verify_aud": False},
        )
        self.assertEqual(decoded["sub"], f"team:{tid}")
        self.assertEqual(decoded["jti"], str(team.active_jti))

    def test_create_sets_request_user_as_owner(self):
        tid = uuid.uuid4()
        self.client.post(
            self.url, {"id": str(tid), "name": "Harper"}, format="json"
        )
        self.assertEqual(Team.objects.get(pk=tid).owner_id, self.user.id)

    def test_same_id_under_other_owner_409s(self):
        tid = uuid.uuid4()
        # Alice creates the team first.
        self.client.post(
            self.url, {"id": str(tid), "name": "Harper"}, format="json"
        )
        # Bob then tries to create with the same id — must be a generic
        # conflict, not idempotent and not 200.
        self.client.force_authenticate(user=self.other_user)
        resp = self.client.post(
            self.url, {"id": str(tid), "name": "Bob's"}, format="json"
        )
        self.assertEqual(resp.status_code, status.HTTP_409_CONFLICT)
        # Owner unchanged.
        self.assertEqual(Team.objects.get(pk=tid).owner_id, self.user.id)

    def test_idempotent_on_same_id_returns_200_without_jwt(self):
        tid = uuid.uuid4()
        first = self.client.post(
            self.url, {"id": str(tid), "name": "first"}, format="json"
        )
        self.assertEqual(first.status_code, status.HTTP_201_CREATED)
        first_jti = Team.objects.get(pk=tid).active_jti

        # Second POST with same id: idempotent hit. Must NOT rotate the
        # jti (otherwise every retry storm re-issues a fresh credential).
        second = self.client.post(
            self.url,
            {"id": str(tid), "name": "ignored-on-hit"},
            format="json",
        )
        self.assertEqual(second.status_code, status.HTTP_200_OK)
        self.assertNotIn("jwt", second.data)
        self.assertEqual(
            Team.objects.get(pk=tid).active_jti, first_jti
        )

    def test_mint_failure_returns_503(self):
        # Retire all signing keys so mint_team_jwt cannot succeed.
        MCPSigningKey.objects.update(is_active=False)
        resp = self.client.post(
            self.url,
            {"id": str(uuid.uuid4()), "name": "x"},
            format="json",
        )
        self.assertEqual(
            resp.status_code, status.HTTP_503_SERVICE_UNAVAILABLE
        )
        # The transaction.atomic() wrapper must also have rolled back the
        # Team row so we don't leave a team with no usable JWT.
        self.assertEqual(Team.objects.count(), 0)


# ---------------------------------------------------------------------------
# GET / DELETE /mcp_server/api/teams/{id}/
# ---------------------------------------------------------------------------


class TeamDetailTest(_AuthenticatedAPITest):
    def setUp(self):
        super().setUp()
        self.team = Team.objects.create(
            id=uuid.uuid4(),
            name="t",
            owner=self.user,
            active=True,
            active_jti=uuid.uuid4(),
        )
        TeamWorkspaceAssignment.objects.create(
            team=self.team, workspace_id="ws-a"
        )
        TeamWorkspaceAssignment.objects.create(
            team=self.team, workspace_id="ws-b"
        )
        self.url = reverse(
            "mcp-server-api:team-detail", kwargs={"team_id": self.team.id}
        )

    def test_get_returns_team_state(self):
        resp = self.client.get(self.url)
        self.assertEqual(resp.status_code, status.HTTP_200_OK)
        self.assertEqual(uuid.UUID(resp.data["id"]), self.team.id)
        self.assertEqual(sorted(resp.data["workspace_ids"]), ["ws-a", "ws-b"])
        # Never leak the JWT via GET.
        self.assertNotIn("jwt", resp.data)

    def test_get_unknown_team_returns_404(self):
        url = reverse(
            "mcp-server-api:team-detail", kwargs={"team_id": uuid.uuid4()}
        )
        resp = self.client.get(url)
        self.assertEqual(resp.status_code, status.HTTP_404_NOT_FOUND)

    def test_delete_soft_deletes(self):
        resp = self.client.delete(self.url)
        self.assertEqual(resp.status_code, status.HTTP_204_NO_CONTENT)

        reloaded = Team.objects.get(pk=self.team.id)
        self.assertFalse(reloaded.active)
        self.assertIsNone(reloaded.active_jti)
        # Workspace rows stay for audit — deactivation only flips flags.
        self.assertEqual(
            TeamWorkspaceAssignment.objects.filter(team=reloaded).count(),
            2,
        )

    def test_delete_idempotent(self):
        self.client.delete(self.url)  # first — 204
        resp = self.client.delete(self.url)
        self.assertEqual(resp.status_code, status.HTTP_204_NO_CONTENT)

    def test_get_by_non_owner_returns_404(self):
        self.client.force_authenticate(user=self.other_user)
        resp = self.client.get(self.url)
        self.assertEqual(resp.status_code, status.HTTP_404_NOT_FOUND)

    def test_delete_by_non_owner_returns_404_and_no_op(self):
        self.client.force_authenticate(user=self.other_user)
        resp = self.client.delete(self.url)
        self.assertEqual(resp.status_code, status.HTTP_404_NOT_FOUND)
        # Original team is still active — non-owner couldn't soft-delete it.
        self.assertTrue(Team.objects.get(pk=self.team.id).active)


# ---------------------------------------------------------------------------
# PUT /mcp_server/api/teams/{id}/workspaces/
# ---------------------------------------------------------------------------


class TeamWorkspacesTest(_AuthenticatedAPITest):
    def setUp(self):
        super().setUp()
        self.team = Team.objects.create(
            id=uuid.uuid4(),
            name="t",
            owner=self.user,
            active=True,
            active_jti=uuid.uuid4(),
        )
        self.url = reverse(
            "mcp-server-api:team-workspaces",
            kwargs={"team_id": self.team.id},
        )

    def _ws_ids(self):
        return sorted(
            self.team.workspace_assignments.values_list(
                "workspace_id", flat=True
            )
        )

    def test_unknown_team_returns_404(self):
        url = reverse(
            "mcp-server-api:team-workspaces",
            kwargs={"team_id": uuid.uuid4()},
        )
        resp = self.client.put(
            url, {"workspace_ids": ["x"]}, format="json"
        )
        self.assertEqual(resp.status_code, status.HTTP_404_NOT_FOUND)

    def test_replace_adds_all(self):
        resp = self.client.put(
            self.url,
            {"workspace_ids": ["ws-a", "ws-b"]},
            format="json",
        )
        self.assertEqual(resp.status_code, status.HTTP_200_OK)
        self.assertEqual(sorted(resp.data["workspace_ids"]), ["ws-a", "ws-b"])
        self.assertEqual(self._ws_ids(), ["ws-a", "ws-b"])

    def test_replace_idempotent_second_call_is_noop(self):
        self.client.put(
            self.url,
            {"workspace_ids": ["ws-a", "ws-b"]},
            format="json",
        )
        existing = list(self.team.workspace_assignments.all())
        pks_before = sorted(a.pk for a in existing)

        # Second identical PUT should not re-create rows.
        resp = self.client.put(
            self.url,
            {"workspace_ids": ["ws-a", "ws-b"]},
            format="json",
        )
        self.assertEqual(resp.status_code, status.HTTP_200_OK)
        pks_after = sorted(
            a.pk for a in self.team.workspace_assignments.all()
        )
        self.assertEqual(pks_before, pks_after)

    def test_replace_removes_dropped(self):
        # Start with a, b, c
        self.client.put(
            self.url,
            {"workspace_ids": ["ws-a", "ws-b", "ws-c"]},
            format="json",
        )
        # Drop b, add d.
        resp = self.client.put(
            self.url,
            {"workspace_ids": ["ws-a", "ws-c", "ws-d"]},
            format="json",
        )
        self.assertEqual(resp.status_code, status.HTTP_200_OK)
        self.assertEqual(self._ws_ids(), ["ws-a", "ws-c", "ws-d"])

    def test_replace_with_empty_set_fail_closed(self):
        self.client.put(
            self.url,
            {"workspace_ids": ["ws-a"]},
            format="json",
        )
        resp = self.client.put(
            self.url, {"workspace_ids": []}, format="json"
        )
        self.assertEqual(resp.status_code, status.HTTP_200_OK)
        self.assertEqual(self._ws_ids(), [])

    def test_duplicates_deduped(self):
        resp = self.client.put(
            self.url,
            {"workspace_ids": ["ws-a", "ws-a", "ws-b"]},
            format="json",
        )
        self.assertEqual(resp.status_code, status.HTTP_200_OK)
        self.assertEqual(sorted(resp.data["workspace_ids"]), ["ws-a", "ws-b"])
        self.assertEqual(self._ws_ids(), ["ws-a", "ws-b"])

    def test_empty_string_rejected(self):
        resp = self.client.put(
            self.url,
            {"workspace_ids": ["ws-a", ""]},
            format="json",
        )
        self.assertEqual(resp.status_code, status.HTTP_400_BAD_REQUEST)

    def test_put_by_non_owner_returns_404(self):
        self.client.force_authenticate(user=self.other_user)
        resp = self.client.put(
            self.url, {"workspace_ids": ["ws-x"]}, format="json"
        )
        self.assertEqual(resp.status_code, status.HTTP_404_NOT_FOUND)
        self.assertEqual(self._ws_ids(), [])


# ---------------------------------------------------------------------------
# POST /mcp_server/api/teams/{id}/rotate/
# ---------------------------------------------------------------------------


class TeamRotateTest(_AuthenticatedAPITest):
    def setUp(self):
        super().setUp()
        _seed_signing_key()
        self.team = Team.objects.create(
            id=uuid.uuid4(),
            name="t",
            owner=self.user,
            active=True,
            active_jti=uuid.uuid4(),
        )
        self.url = reverse(
            "mcp-server-api:team-rotate",
            kwargs={"team_id": self.team.id},
        )

    def test_rotate_upserts_missing_team(self):
        # Rotate is upsert-on-missing: if no Team row exists for this
        # id, create one owned by the caller and mint its first JWT.
        # Eliminates the create-before-rotate ordering trap Daedalus hit
        # in production.
        new_id = uuid.uuid4()
        url = reverse(
            "mcp-server-api:team-rotate",
            kwargs={"team_id": new_id},
        )
        resp = self.client.post(url)
        self.assertEqual(resp.status_code, status.HTTP_200_OK)
        self.assertIn("jwt", resp.data)
        team = Team.objects.get(pk=new_id)
        self.assertEqual(team.owner_id, self.user.id)
        self.assertTrue(team.active)
        self.assertIsNotNone(team.active_jti)

    def test_rotate_returns_new_jwt_and_changes_active_jti(self):
        before = self.team.active_jti
        resp = self.client.post(self.url)
        self.assertEqual(resp.status_code, status.HTTP_200_OK)
        self.assertIn("jwt", resp.data)
        self.team.refresh_from_db()
        self.assertNotEqual(self.team.active_jti, before)

    def test_rotate_inactive_team_409(self):
        self.team.deactivate()
        resp = self.client.post(self.url)
        self.assertEqual(resp.status_code, status.HTTP_409_CONFLICT)
        # And we did NOT revive the team by side effect.
        self.team.refresh_from_db()
        self.assertFalse(self.team.active)
        self.assertIsNone(self.team.active_jti)

    def test_rotate_without_signing_key_returns_503(self):
        MCPSigningKey.objects.update(is_active=False)
        resp = self.client.post(self.url)
        self.assertEqual(
            resp.status_code, status.HTTP_503_SERVICE_UNAVAILABLE
        )

    def test_rotate_by_non_owner_returns_409(self):
        # The team row exists under Alice; Bob rotating it must not
        # upsert (that would silently steal the id) and must not 404
        # (would tell Bob the id is free). 409 is the right answer.
        before = self.team.active_jti
        before_owner = self.team.owner_id
        self.client.force_authenticate(user=self.other_user)
        resp = self.client.post(self.url)
        self.assertEqual(resp.status_code, status.HTTP_409_CONFLICT)
        # Alice's team is untouched.
        self.team.refresh_from_db()
        self.assertEqual(self.team.active_jti, before)
        self.assertEqual(self.team.owner_id, before_owner)