test(mcp): RevokeBySession + ReapStaleSystemTokens integration

This commit is contained in:
2026-05-08 07:27:59 +08:00
parent d431b8103d
commit 4071778930
+117
View File
@@ -271,3 +271,120 @@ func TestPgRepo_PurgeExpired(t *testing.T) {
require.True(t, ok) require.True(t, ok)
assert.Equal(t, errs.CodeMcpTokenNotFound, ae.Code) assert.Equal(t, errs.CodeMcpTokenNotFound, ae.Code)
} }
// mustInsertACPSetup 准备测 system token 需要的 fixtures:
// project + workspace + agent_kind + acp_sessions row(status 任选)。
// 返回 sessionID + projectID(用于 system token 的 ACPSessionID + scope.project_ids)。
func mustInsertACPSetup(t *testing.T, ctx context.Context, pool *pgxpool.Pool, ownerUID uuid.UUID, status string) (sessionID, projectID uuid.UUID) {
t.Helper()
// project
pid := uuid.New()
_, err := pool.Exec(ctx, `
INSERT INTO projects (id, slug, name, owner_id, visibility)
VALUES ($1, $2, $2, $3, 'private')`,
pid, "p"+pid.String()[:8], ownerUID)
require.NoError(t, err)
// workspace
wid := uuid.New()
_, err = pool.Exec(ctx, `
INSERT INTO workspaces (id, project_id, slug, name, git_remote_url, default_branch, main_path, sync_status)
VALUES ($1, $2, $3, $3, 'git@local:x.git', 'main', '/tmp/'||$1::text, 'idle')`,
wid, pid, "w"+wid.String()[:8])
require.NoError(t, err)
// agent_kind
akid := uuid.New()
_, err = pool.Exec(ctx, `
INSERT INTO acp_agent_kinds (id, name, display_name, binary_path, args, enabled, created_by)
VALUES ($1, $2, $2, '/bin/echo', '{}', TRUE, $3)`,
akid, "ak"+akid.String()[:8], ownerUID)
require.NoError(t, err)
// session
sid := uuid.New()
_, err = pool.Exec(ctx, `
INSERT INTO acp_sessions (id, workspace_id, project_id, agent_kind_id, user_id, branch, cwd_path, is_main_worktree, status)
VALUES ($1, $2, $3, $4, $5, 'main', '/tmp/x', TRUE, $6)`,
sid, wid, pid, akid, ownerUID, status)
require.NoError(t, err)
return sid, pid
}
func TestPgRepo_RevokeBySession(t *testing.T) {
t.Parallel()
ctx := context.Background()
repo, pool := setupRepo(t)
uid := mustInsertUser(t, ctx, pool, "rs@local", false)
sid, pid := mustInsertACPSetup(t, ctx, pool, uid, "running")
// 插一个 system token
_, hash, _ := mcp.NewToken()
tk, err := repo.Create(ctx, &mcp.MCPToken{
ID: uuid.New(), TokenHash: hash, UserID: uid,
Name: "system-x", Issuer: mcp.IssuerSystem,
Scope: mcp.Scope{ProjectIDs: []uuid.UUID{pid}},
ACPSessionID: &sid,
CreatedBy: uid,
})
require.NoError(t, err)
// 撤销
revoked, err := repo.RevokeBySession(ctx, sid)
require.NoError(t, err)
require.Len(t, revoked, 1)
assert.Equal(t, tk.ID, revoked[0])
// row 上 revoked_at 写入
got, err := repo.GetByID(ctx, tk.ID)
require.NoError(t, err)
require.NotNil(t, got.RevokedAt)
// 重复调用 → 0 行(已撤销不再返)
again, err := repo.RevokeBySession(ctx, sid)
require.NoError(t, err)
assert.Empty(t, again)
}
func TestPgRepo_ReapStaleSystemTokens(t *testing.T) {
t.Parallel()
ctx := context.Background()
repo, pool := setupRepo(t)
uid := mustInsertUser(t, ctx, pool, "reap@local", false)
// 一个 crashed session 的 system token
sidCrashed, pidA := mustInsertACPSetup(t, ctx, pool, uid, "crashed")
_, h1, _ := mcp.NewToken()
tk1, err := repo.Create(ctx, &mcp.MCPToken{
ID: uuid.New(), TokenHash: h1, UserID: uid,
Name: "crashed-tok", Issuer: mcp.IssuerSystem,
Scope: mcp.Scope{ProjectIDs: []uuid.UUID{pidA}},
ACPSessionID: &sidCrashed, CreatedBy: uid,
})
require.NoError(t, err)
// 一个 running session 的 system token(不该被 reap)
sidRunning, pidB := mustInsertACPSetup(t, ctx, pool, uid, "running")
_, h2, _ := mcp.NewToken()
tk2, err := repo.Create(ctx, &mcp.MCPToken{
ID: uuid.New(), TokenHash: h2, UserID: uid,
Name: "running-tok", Issuer: mcp.IssuerSystem,
Scope: mcp.Scope{ProjectIDs: []uuid.UUID{pidB}},
ACPSessionID: &sidRunning, CreatedBy: uid,
})
require.NoError(t, err)
reaped, err := repo.ReapStaleSystemTokens(ctx)
require.NoError(t, err)
require.Len(t, reaped, 1)
assert.Equal(t, tk1.ID, reaped[0])
got1, _ := repo.GetByID(ctx, tk1.ID)
require.NotNil(t, got1.RevokedAt)
got2, _ := repo.GetByID(ctx, tk2.ID)
assert.Nil(t, got2.RevokedAt)
}