test(mcp): repository CRUD + CHECK constraints + filter + purge integration

2026-05-08 07:25:21 +08:00
parent bb512c0029
commit d431b8103d
1 changed files with 273 additions and 0 deletions
@@ -0,0 +1,273 @@
+//go:build integration
+
+package mcp_test
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/jackc/pgx/v5/pgxpool"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	tcpg "github.com/testcontainers/testcontainers-go/modules/postgres"
+
+	"github.com/yan1h/agent-coding-workflow/internal/infra/db"
+	"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
+	"github.com/yan1h/agent-coding-workflow/internal/mcp"
+)
+
+func setupRepo(t *testing.T) (mcp.Repository, *pgxpool.Pool) {
+	t.Helper()
+	ctx := context.Background()
+	container, err := tcpg.Run(ctx, "postgres:16-alpine",
+		tcpg.WithDatabase("acw"),
+		tcpg.WithUsername("acw"),
+		tcpg.WithPassword("acw"),
+	)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = container.Terminate(ctx) })
+
+	dsn, err := container.ConnectionString(ctx, "sslmode=disable")
+	require.NoError(t, err)
+	require.NoError(t, db.Migrate(ctx, dsn, "file://../../migrations"))
+
+	pool, err := db.NewPool(ctx, dsn)
+	require.NoError(t, err)
+	t.Cleanup(pool.Close)
+
+	return mcp.NewPostgresRepository(pool), pool
+}
+
+// mustInsertUser 直插一条 user，返 UID。沿用 acp 测试模式。
+func mustInsertUser(t *testing.T, ctx context.Context, pool *pgxpool.Pool, email string, isAdmin bool) uuid.UUID {
+	t.Helper()
+	uid := uuid.New()
+	_, err := pool.Exec(ctx, `
+		INSERT INTO users (id, email, password_hash, display_name, is_admin)
+		VALUES ($1, $2, 'argon2id$dummy', $3, $4)`,
+		uid, email, email, isAdmin)
+	require.NoError(t, err)
+	return uid
+}
+
+func TestPgRepo_AdminToken_CRUD(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	repo, pool := setupRepo(t)
+
+	uid := mustInsertUser(t, ctx, pool, "admin1@local", true)
+	tok, hash, err := mcp.NewToken()
+	require.NoError(t, err)
+
+	expires := time.Now().Add(90 * 24 * time.Hour)
+	created, err := repo.Create(ctx, &mcp.MCPToken{
+		ID:        uuid.New(),
+		TokenHash: hash,
+		UserID:    uid,
+		Name:      "test admin token",
+		Issuer:    mcp.IssuerAdmin,
+		Scope:     mcp.Scope{Tools: []string{"list_issues"}, ProjectIDs: nil},
+		ExpiresAt: &expires,
+		CreatedBy: uid,
+	})
+	require.NoError(t, err)
+	assert.NotEqual(t, uuid.Nil, created.ID)
+	assert.Equal(t, mcp.IssuerAdmin, created.Issuer)
+	assert.Equal(t, []string{"list_issues"}, created.Scope.Tools)
+
+	// GetByHash
+	got, err := repo.GetByHash(ctx, mcp.HashToken(tok))
+	require.NoError(t, err)
+	assert.Equal(t, created.ID, got.ID)
+
+	// GetByID
+	got2, err := repo.GetByID(ctx, created.ID)
+	require.NoError(t, err)
+	assert.Equal(t, created.ID, got2.ID)
+
+	// List (admin call: callerUserID nil)
+	all, err := repo.List(ctx, nil, false)
+	require.NoError(t, err)
+	require.Len(t, all, 1)
+	assert.Equal(t, created.ID, all[0].ID)
+
+	// UpdateLastUsed
+	require.NoError(t, repo.UpdateLastUsed(ctx, created.ID))
+	got3, err := repo.GetByID(ctx, created.ID)
+	require.NoError(t, err)
+	require.NotNil(t, got3.LastUsedAt)
+	assert.WithinDuration(t, time.Now(), *got3.LastUsedAt, 5*time.Second)
+
+	// Revoke
+	revokedID, err := repo.Revoke(ctx, created.ID)
+	require.NoError(t, err)
+	assert.Equal(t, created.ID, revokedID)
+
+	// 二次 Revoke → NotFound
+	_, err = repo.Revoke(ctx, created.ID)
+	ae, ok := errs.As(err)
+	require.True(t, ok)
+	assert.Equal(t, errs.CodeMcpTokenNotFound, ae.Code)
+}
+
+func TestPgRepo_GetByHash_NotFound(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	repo, _ := setupRepo(t)
+
+	_, err := repo.GetByHash(ctx, mcp.HashToken("mcpt_nope"))
+	ae, ok := errs.As(err)
+	require.True(t, ok)
+	assert.Equal(t, errs.CodeMcpTokenInvalid, ae.Code)
+}
+
+func TestPgRepo_CheckConstraint_SystemRequiresSession(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	repo, pool := setupRepo(t)
+
+	uid := mustInsertUser(t, ctx, pool, "u@local", false)
+	_, hash, err := mcp.NewToken()
+	require.NoError(t, err)
+
+	// system token + 没 acp_session_id → CHECK 拒绝
+	_, err = repo.Create(ctx, &mcp.MCPToken{
+		ID: uuid.New(), TokenHash: hash, UserID: uid,
+		Name: "x", Issuer: mcp.IssuerSystem, CreatedBy: uid,
+	})
+	require.Error(t, err)
+	// PG 抛 23514 check_violation；包装在 errs.CodeInternal 里
+}
+
+func TestPgRepo_CheckConstraint_AdminMustNotHaveSession(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	repo, pool := setupRepo(t)
+
+	uid := mustInsertUser(t, ctx, pool, "u2@local", false)
+	_, hash, err := mcp.NewToken()
+	require.NoError(t, err)
+
+	fakeSessionID := uuid.New()
+	// admin token + 有 acp_session_id → CHECK 拒绝
+	_, err = repo.Create(ctx, &mcp.MCPToken{
+		ID: uuid.New(), TokenHash: hash, UserID: uid,
+		Name: "y", Issuer: mcp.IssuerAdmin, ACPSessionID: &fakeSessionID,
+		CreatedBy: uid,
+	})
+	require.Error(t, err)
+}
+
+func TestPgRepo_List_FilterByCaller(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	repo, pool := setupRepo(t)
+
+	admin := mustInsertUser(t, ctx, pool, "admin@local", true)
+	user1 := mustInsertUser(t, ctx, pool, "u1@local", false)
+	user2 := mustInsertUser(t, ctx, pool, "u2@local", false)
+
+	mkAdmin := func(uid uuid.UUID, name string) {
+		_, hash, _ := mcp.NewToken()
+		_, err := repo.Create(ctx, &mcp.MCPToken{
+			ID: uuid.New(), TokenHash: hash, UserID: uid, Name: name,
+			Issuer: mcp.IssuerAdmin, CreatedBy: admin,
+		})
+		require.NoError(t, err)
+	}
+	mkAdmin(user1, "u1.t1")
+	mkAdmin(user1, "u1.t2")
+	mkAdmin(user2, "u2.t1")
+
+	// admin 视角：看全部
+	all, err := repo.List(ctx, nil, false)
+	require.NoError(t, err)
+	assert.Len(t, all, 3)
+
+	// user1 视角：仅自己持有的 admin-issued
+	mine, err := repo.List(ctx, &user1, false)
+	require.NoError(t, err)
+	assert.Len(t, mine, 2)
+	for _, t2 := range mine {
+		assert.Equal(t, user1, t2.UserID)
+	}
+}
+
+func TestPgRepo_List_ActiveOnly(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	repo, pool := setupRepo(t)
+
+	admin := mustInsertUser(t, ctx, pool, "a@local", true)
+	user := mustInsertUser(t, ctx, pool, "u@local", false)
+
+	// 一个有效，一个已撤销，一个已过期
+	expFuture := time.Now().Add(24 * time.Hour)
+	expPast := time.Now().Add(-1 * time.Hour)
+	mk := func(name string, exp *time.Time) uuid.UUID {
+		_, hash, _ := mcp.NewToken()
+		tok, err := repo.Create(ctx, &mcp.MCPToken{
+			ID: uuid.New(), TokenHash: hash, UserID: user, Name: name,
+			Issuer: mcp.IssuerAdmin, ExpiresAt: exp, CreatedBy: admin,
+		})
+		require.NoError(t, err)
+		return tok.ID
+	}
+	mk("active", &expFuture)
+	revID := mk("revoked", &expFuture)
+	mk("expired", &expPast)
+
+	_, err := repo.Revoke(ctx, revID)
+	require.NoError(t, err)
+
+	all, err := repo.List(ctx, nil, false)
+	require.NoError(t, err)
+	assert.Len(t, all, 3)
+
+	active, err := repo.List(ctx, nil, true)
+	require.NoError(t, err)
+	require.Len(t, active, 1)
+	assert.Equal(t, "active", active[0].Name)
+}
+
+func TestPgRepo_PurgeExpired(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	repo, pool := setupRepo(t)
+
+	admin := mustInsertUser(t, ctx, pool, "p@local", true)
+	user := mustInsertUser(t, ctx, pool, "p2@local", false)
+
+	old := time.Now().Add(-30 * 24 * time.Hour)
+	expFuture := time.Now().Add(24 * time.Hour)
+
+	// 一个早已过期的 row（手工写 expires_at = 30d ago）
+	_, hash, _ := mcp.NewToken()
+	tk, err := repo.Create(ctx, &mcp.MCPToken{
+		ID: uuid.New(), TokenHash: hash, UserID: user, Name: "old",
+		Issuer: mcp.IssuerAdmin, ExpiresAt: &old, CreatedBy: admin,
+	})
+	require.NoError(t, err)
+
+	// 一个新 token
+	_, hash2, _ := mcp.NewToken()
+	_, err = repo.Create(ctx, &mcp.MCPToken{
+		ID: uuid.New(), TokenHash: hash2, UserID: user, Name: "new",
+		Issuer: mcp.IssuerAdmin, ExpiresAt: &expFuture, CreatedBy: admin,
+	})
+	require.NoError(t, err)
+
+	// 删 7 天前 retention
+	cutoff := time.Now().Add(-7 * 24 * time.Hour)
+	n, err := repo.PurgeExpired(ctx, cutoff)
+	require.NoError(t, err)
+	assert.Equal(t, int64(1), n)
+
+	// 老 token 应已删
+	_, err = repo.GetByID(ctx, tk.ID)
+	ae, ok := errs.As(err)
+	require.True(t, ok)
+	assert.Equal(t, errs.CodeMcpTokenNotFound, ae.Code)
+}