You've already forked agentic-coding-workflow
274 lines
7.7 KiB
Go
274 lines
7.7 KiB
Go
//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)
|
|
}
|