diff --git a/internal/mcp/repository_test.go b/internal/mcp/repository_test.go new file mode 100644 index 0000000..f841dff --- /dev/null +++ b/internal/mcp/repository_test.go @@ -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) +}