You've already forked agentic-coding-workflow
test(mcp): TokenService unit tests with fake repo
This commit is contained in:
@@ -0,0 +1,341 @@
|
||||
package mcp_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/yan1h/agent-coding-workflow/internal/audit"
|
||||
"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
|
||||
"github.com/yan1h/agent-coding-workflow/internal/mcp"
|
||||
)
|
||||
|
||||
// ===== fake repo =====
|
||||
|
||||
type fakeRepo struct {
|
||||
mu sync.Mutex
|
||||
tokens map[uuid.UUID]*mcp.MCPToken
|
||||
byHash map[string]uuid.UUID
|
||||
}
|
||||
|
||||
func newFakeRepo() *fakeRepo {
|
||||
return &fakeRepo{
|
||||
tokens: map[uuid.UUID]*mcp.MCPToken{},
|
||||
byHash: map[string]uuid.UUID{},
|
||||
}
|
||||
}
|
||||
|
||||
func (r *fakeRepo) Create(ctx context.Context, t *mcp.MCPToken) (*mcp.MCPToken, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
c := *t
|
||||
c.CreatedAt = time.Now()
|
||||
r.tokens[c.ID] = &c
|
||||
r.byHash[string(c.TokenHash)] = c.ID
|
||||
return &c, nil
|
||||
}
|
||||
func (r *fakeRepo) GetByHash(ctx context.Context, h []byte) (*mcp.MCPToken, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
id, ok := r.byHash[string(h)]
|
||||
if !ok {
|
||||
return nil, errs.New(errs.CodeMcpTokenInvalid, "not found")
|
||||
}
|
||||
c := *r.tokens[id]
|
||||
return &c, nil
|
||||
}
|
||||
func (r *fakeRepo) GetByID(ctx context.Context, id uuid.UUID) (*mcp.MCPToken, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
t, ok := r.tokens[id]
|
||||
if !ok {
|
||||
return nil, errs.New(errs.CodeMcpTokenNotFound, "not found")
|
||||
}
|
||||
c := *t
|
||||
return &c, nil
|
||||
}
|
||||
func (r *fakeRepo) List(ctx context.Context, caller *uuid.UUID, activeOnly bool) ([]*mcp.MCPToken, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
var out []*mcp.MCPToken
|
||||
for _, t := range r.tokens {
|
||||
if caller != nil && (t.UserID != *caller || t.Issuer != mcp.IssuerAdmin) {
|
||||
continue
|
||||
}
|
||||
if activeOnly && !t.IsActive(time.Now()) {
|
||||
continue
|
||||
}
|
||||
c := *t
|
||||
out = append(out, &c)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
func (r *fakeRepo) UpdateLastUsed(ctx context.Context, id uuid.UUID) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if t, ok := r.tokens[id]; ok {
|
||||
now := time.Now()
|
||||
t.LastUsedAt = &now
|
||||
}
|
||||
return nil
|
||||
}
|
||||
func (r *fakeRepo) Revoke(ctx context.Context, id uuid.UUID) (uuid.UUID, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
t, ok := r.tokens[id]
|
||||
if !ok || t.RevokedAt != nil {
|
||||
return uuid.Nil, errs.New(errs.CodeMcpTokenNotFound, "not found")
|
||||
}
|
||||
now := time.Now()
|
||||
t.RevokedAt = &now
|
||||
return id, nil
|
||||
}
|
||||
func (r *fakeRepo) RevokeBySession(ctx context.Context, sid uuid.UUID) ([]uuid.UUID, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
var revoked []uuid.UUID
|
||||
now := time.Now()
|
||||
for _, t := range r.tokens {
|
||||
if t.Issuer == mcp.IssuerSystem && t.ACPSessionID != nil && *t.ACPSessionID == sid && t.RevokedAt == nil {
|
||||
t.RevokedAt = &now
|
||||
revoked = append(revoked, t.ID)
|
||||
}
|
||||
}
|
||||
return revoked, nil
|
||||
}
|
||||
func (r *fakeRepo) ReapStaleSystemTokens(ctx context.Context) ([]uuid.UUID, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (r *fakeRepo) PurgeExpired(ctx context.Context, before time.Time) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
func (r *fakeRepo) DeleteByID(ctx context.Context, id uuid.UUID) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
delete(r.tokens, id)
|
||||
return nil
|
||||
}
|
||||
|
||||
// ===== fake audit =====
|
||||
|
||||
type fakeAudit struct {
|
||||
mu sync.Mutex
|
||||
entries []audit.Entry
|
||||
}
|
||||
|
||||
func (a *fakeAudit) Record(ctx context.Context, e audit.Entry) error {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
a.entries = append(a.entries, e)
|
||||
return nil
|
||||
}
|
||||
func (a *fakeAudit) actions() []string {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
out := make([]string, 0, len(a.entries))
|
||||
for _, e := range a.entries {
|
||||
out = append(out, e.Action)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// ===== tests =====
|
||||
|
||||
func TestIssueAdmin_HappyPath(t *testing.T) {
|
||||
t.Parallel()
|
||||
repo := newFakeRepo()
|
||||
au := &fakeAudit{}
|
||||
svc := mcp.NewTokenService(repo, au)
|
||||
|
||||
uid := uuid.New()
|
||||
in := mcp.IssueAdminInput{
|
||||
UserID: uid,
|
||||
Name: "test",
|
||||
Scope: mcp.Scope{Tools: []string{"list_issues"}},
|
||||
ExpiresAt: time.Now().Add(time.Hour),
|
||||
CreatedBy: uid,
|
||||
}
|
||||
res, err := svc.IssueAdmin(context.Background(), in)
|
||||
require.NoError(t, err)
|
||||
assert.NotEmpty(t, res.Plaintext)
|
||||
assert.Equal(t, []string{"mcp.token.create"}, au.actions())
|
||||
}
|
||||
|
||||
func TestIssueAdmin_NameRequired(t *testing.T) {
|
||||
t.Parallel()
|
||||
svc := mcp.NewTokenService(newFakeRepo(), &fakeAudit{})
|
||||
|
||||
_, err := svc.IssueAdmin(context.Background(), mcp.IssueAdminInput{
|
||||
UserID: uuid.New(),
|
||||
Name: "",
|
||||
ExpiresAt: time.Now().Add(time.Hour),
|
||||
CreatedBy: uuid.New(),
|
||||
})
|
||||
ae, ok := errs.As(err)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, errs.CodeMcpTokenNameRequired, ae.Code)
|
||||
}
|
||||
|
||||
func TestIssueAdmin_ExpiresRequired(t *testing.T) {
|
||||
t.Parallel()
|
||||
svc := mcp.NewTokenService(newFakeRepo(), &fakeAudit{})
|
||||
|
||||
_, err := svc.IssueAdmin(context.Background(), mcp.IssueAdminInput{
|
||||
UserID: uuid.New(),
|
||||
Name: "x",
|
||||
CreatedBy: uuid.New(),
|
||||
})
|
||||
ae, ok := errs.As(err)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, errs.CodeMcpTokenExpiresRequired, ae.Code)
|
||||
}
|
||||
|
||||
func TestIssueSystemToken_Defaults(t *testing.T) {
|
||||
t.Parallel()
|
||||
repo := newFakeRepo()
|
||||
au := &fakeAudit{}
|
||||
svc := mcp.NewTokenService(repo, au)
|
||||
|
||||
uid := uuid.New()
|
||||
sid := uuid.New()
|
||||
pid := uuid.New()
|
||||
res, err := svc.IssueSystemToken(context.Background(), mcp.IssueSystemTokenInput{
|
||||
UserID: uid,
|
||||
ACPSessionID: sid,
|
||||
ProjectIDs: []uuid.UUID{pid},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
got, err := repo.GetByID(context.Background(), res.ID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, mcp.IssuerSystem, got.Issuer)
|
||||
require.NotNil(t, got.ACPSessionID)
|
||||
assert.Equal(t, sid, *got.ACPSessionID)
|
||||
assert.NotNil(t, got.ExpiresAt)
|
||||
assert.WithinDuration(t, time.Now().Add(24*time.Hour), *got.ExpiresAt, time.Minute)
|
||||
// 默认工具白名单包含 PM 写入但不含 list_recent_messages
|
||||
assert.Contains(t, got.Scope.Tools, "create_issue")
|
||||
assert.NotContains(t, got.Scope.Tools, "list_recent_messages")
|
||||
assert.Equal(t, []string{"mcp.token.create_system"}, au.actions())
|
||||
}
|
||||
|
||||
func TestAuthenticate_HappyPath(t *testing.T) {
|
||||
t.Parallel()
|
||||
repo := newFakeRepo()
|
||||
svc := mcp.NewTokenService(repo, &fakeAudit{})
|
||||
|
||||
res, err := svc.IssueAdmin(context.Background(), mcp.IssueAdminInput{
|
||||
UserID: uuid.New(),
|
||||
Name: "x",
|
||||
ExpiresAt: time.Now().Add(time.Hour),
|
||||
CreatedBy: uuid.New(),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
tok, err := svc.Authenticate(context.Background(), res.Plaintext)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, res.ID, tok.ID)
|
||||
}
|
||||
|
||||
func TestAuthenticate_Empty(t *testing.T) {
|
||||
t.Parallel()
|
||||
svc := mcp.NewTokenService(newFakeRepo(), &fakeAudit{})
|
||||
|
||||
_, err := svc.Authenticate(context.Background(), "")
|
||||
ae, ok := errs.As(err)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, errs.CodeMcpTokenInvalid, ae.Code)
|
||||
}
|
||||
|
||||
func TestAuthenticate_Revoked(t *testing.T) {
|
||||
t.Parallel()
|
||||
repo := newFakeRepo()
|
||||
svc := mcp.NewTokenService(repo, &fakeAudit{})
|
||||
|
||||
res, _ := svc.IssueAdmin(context.Background(), mcp.IssueAdminInput{
|
||||
UserID: uuid.New(),
|
||||
Name: "x",
|
||||
ExpiresAt: time.Now().Add(time.Hour),
|
||||
CreatedBy: uuid.New(),
|
||||
})
|
||||
require.NoError(t, svc.Revoke(context.Background(), res.ID, uuid.New()))
|
||||
|
||||
_, err := svc.Authenticate(context.Background(), res.Plaintext)
|
||||
ae, ok := errs.As(err)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, errs.CodeMcpTokenRevoked, ae.Code)
|
||||
}
|
||||
|
||||
func TestAuthenticate_Expired(t *testing.T) {
|
||||
t.Parallel()
|
||||
repo := newFakeRepo()
|
||||
svc := mcp.NewTokenService(repo, &fakeAudit{})
|
||||
|
||||
res, _ := svc.IssueAdmin(context.Background(), mcp.IssueAdminInput{
|
||||
UserID: uuid.New(),
|
||||
Name: "x",
|
||||
ExpiresAt: time.Now().Add(time.Hour),
|
||||
CreatedBy: uuid.New(),
|
||||
})
|
||||
// 直接改 fake repo 中 expires_at 为已过去
|
||||
past := time.Now().Add(-time.Hour)
|
||||
repo.tokens[res.ID].ExpiresAt = &past
|
||||
|
||||
_, err := svc.Authenticate(context.Background(), res.Plaintext)
|
||||
ae, ok := errs.As(err)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, errs.CodeMcpTokenExpired, ae.Code)
|
||||
}
|
||||
|
||||
func TestAuthenticate_TamperedToken(t *testing.T) {
|
||||
t.Parallel()
|
||||
svc := mcp.NewTokenService(newFakeRepo(), &fakeAudit{})
|
||||
|
||||
_, err := svc.Authenticate(context.Background(), "mcpt_tampered_value")
|
||||
ae, ok := errs.As(err)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, errs.CodeMcpTokenInvalid, ae.Code)
|
||||
}
|
||||
|
||||
func TestRevokeBySession_AuditPerToken(t *testing.T) {
|
||||
t.Parallel()
|
||||
repo := newFakeRepo()
|
||||
au := &fakeAudit{}
|
||||
svc := mcp.NewTokenService(repo, au)
|
||||
|
||||
uid := uuid.New()
|
||||
sid := uuid.New()
|
||||
pid := uuid.New()
|
||||
for i := 0; i < 3; i++ {
|
||||
_, err := svc.IssueSystemToken(context.Background(), mcp.IssueSystemTokenInput{
|
||||
UserID: uid, ACPSessionID: sid, ProjectIDs: []uuid.UUID{pid},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
require.NoError(t, svc.RevokeBySession(context.Background(), sid))
|
||||
|
||||
// 3 个 system 创建 + 3 个 by_session 撤销 = 6 条 audit
|
||||
actions := au.actions()
|
||||
assert.Len(t, actions, 6)
|
||||
createCnt, revokeCnt := 0, 0
|
||||
for _, a := range actions {
|
||||
switch a {
|
||||
case "mcp.token.create_system":
|
||||
createCnt++
|
||||
case "mcp.token.revoke_by_session":
|
||||
revokeCnt++
|
||||
}
|
||||
}
|
||||
assert.Equal(t, 3, createCnt)
|
||||
assert.Equal(t, 3, revokeCnt)
|
||||
}
|
||||
|
||||
// 编译期保证：类型实现接口
|
||||
var _ = errors.New
|
||||
Reference in New Issue
Block a user