diff --git a/internal/acp/repository_test.go b/internal/acp/repository_test.go new file mode 100644 index 0000000..fcaf2d4 --- /dev/null +++ b/internal/acp/repository_test.go @@ -0,0 +1,261 @@ +//go:build integration + +package acp_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/acp" + "github.com/yan1h/agent-coding-workflow/internal/infra/db" + "github.com/yan1h/agent-coding-workflow/internal/infra/errs" +) + +func setupRepo(t *testing.T) (acp.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 acp.NewPostgresRepository(pool), pool +} + +func TestPgRepo_AgentKind_CRUD(t *testing.T) { + t.Parallel() + ctx := context.Background() + repo, pool := setupRepo(t) + + uid := mustInsertUser(t, ctx, pool, "admin@local", true) + + // Create + k, err := repo.CreateAgentKind(ctx, &acp.AgentKind{ + ID: uuid.New(), Name: "claude_code", DisplayName: "Claude Code", + BinaryPath: "claude-code", Args: []string{"--acp"}, + EncryptedEnv: []byte("ciphertext"), Enabled: true, CreatedBy: uid, + }) + require.NoError(t, err) + assert.Equal(t, "claude_code", k.Name) + + // Get by ID + got, err := repo.GetAgentKindByID(ctx, k.ID) + require.NoError(t, err) + assert.Equal(t, k.Name, got.Name) + + // Get by Name + got2, err := repo.GetAgentKindByName(ctx, "claude_code") + require.NoError(t, err) + assert.Equal(t, k.ID, got2.ID) + + // List enabled + enabled, err := repo.ListEnabledAgentKinds(ctx) + require.NoError(t, err) + require.Len(t, enabled, 1) + + // Update(enabled=false) + k.DisplayName = "Claude Code v2" + k.Enabled = false + updated, err := repo.UpdateAgentKind(ctx, k) + require.NoError(t, err) + assert.Equal(t, "Claude Code v2", updated.DisplayName) + assert.False(t, updated.Enabled) + + enabled, _ = repo.ListEnabledAgentKinds(ctx) + assert.Empty(t, enabled) + + // Delete + require.NoError(t, repo.DeleteAgentKind(ctx, k.ID)) + _, err = repo.GetAgentKindByID(ctx, k.ID) + ae, ok := errs.As(err) + require.True(t, ok) + assert.Equal(t, errs.CodeAcpAgentKindNotFound, ae.Code) +} + +func TestPgRepo_AgentKind_DeleteInUse(t *testing.T) { + t.Parallel() + ctx := context.Background() + repo, pool := setupRepo(t) + + uid := mustInsertUser(t, ctx, pool, "admin2@local", true) + pid, wsid := mustInsertProjectWorkspace(t, ctx, pool, uid) + + k, err := repo.CreateAgentKind(ctx, &acp.AgentKind{ + ID: uuid.New(), Name: "kind-in-use", DisplayName: "x", + BinaryPath: "x", Enabled: true, CreatedBy: uid, + }) + require.NoError(t, err) + + _, err = repo.InsertSession(ctx, &acp.Session{ + ID: uuid.New(), WorkspaceID: wsid, ProjectID: pid, + AgentKindID: k.ID, UserID: uid, + Branch: "main", CwdPath: "/tmp", IsMainWorktree: true, Status: acp.SessionRunning, + }) + require.NoError(t, err) + + err = repo.DeleteAgentKind(ctx, k.ID) + ae, ok := errs.As(err) + require.True(t, ok) + assert.Equal(t, errs.CodeAcpAgentKindInUse, ae.Code) +} + +func TestPgRepo_Session_Lifecycle(t *testing.T) { + t.Parallel() + ctx := context.Background() + repo, pool := setupRepo(t) + + uid := mustInsertUser(t, ctx, pool, "u3@local", false) + pid, wsid := mustInsertProjectWorkspace(t, ctx, pool, uid) + k, _ := repo.CreateAgentKind(ctx, &acp.AgentKind{ + ID: uuid.New(), Name: "kind3", DisplayName: "x", + BinaryPath: "x", Enabled: true, CreatedBy: uid, + }) + + sess, err := repo.InsertSession(ctx, &acp.Session{ + ID: uuid.New(), WorkspaceID: wsid, ProjectID: pid, + AgentKindID: k.ID, UserID: uid, + Branch: "feat-x", CwdPath: "/tmp", IsMainWorktree: false, Status: acp.SessionStarting, + }) + require.NoError(t, err) + + // Count active + n, _ := repo.CountActiveSessions(ctx) + assert.Equal(t, int64(1), n) + n2, _ := repo.CountActiveSessionsByUser(ctx, uid) + assert.Equal(t, int64(1), n2) + + // Update running + require.NoError(t, repo.UpdateSessionRunning(ctx, sess.ID, "agent-sid-abc", 12345)) + got, _ := repo.GetSessionByID(ctx, sess.ID) + assert.Equal(t, acp.SessionRunning, got.Status) + require.NotNil(t, got.AgentSessionID) + assert.Equal(t, "agent-sid-abc", *got.AgentSessionID) + require.NotNil(t, got.PID) + assert.Equal(t, int32(12345), *got.PID) + + // Update finished + ec := int32(0) + require.NoError(t, repo.UpdateSessionFinished(ctx, sess.ID, acp.SessionExited, &ec, nil)) + got, _ = repo.GetSessionByID(ctx, sess.ID) + assert.Equal(t, acp.SessionExited, got.Status) + assert.False(t, got.Status.IsActive()) + + // Reset stuck(应无影响:本 session 已 exited) + rows, _ := repo.ResetStuckSessionsOnRestart(ctx) + assert.Empty(t, rows) +} + +func TestPgRepo_Session_ResetStuck(t *testing.T) { + t.Parallel() + ctx := context.Background() + repo, pool := setupRepo(t) + + uid := mustInsertUser(t, ctx, pool, "u4@local", false) + pid, wsid := mustInsertProjectWorkspace(t, ctx, pool, uid) + k, _ := repo.CreateAgentKind(ctx, &acp.AgentKind{ + ID: uuid.New(), Name: "kind4", DisplayName: "x", + BinaryPath: "x", Enabled: true, CreatedBy: uid, + }) + + for i := 0; i < 2; i++ { + _, _ = repo.InsertSession(ctx, &acp.Session{ + ID: uuid.New(), WorkspaceID: wsid, ProjectID: pid, + AgentKindID: k.ID, UserID: uid, + Branch: "stuck", CwdPath: "/tmp", IsMainWorktree: false, Status: acp.SessionRunning, + }) + } + + rows, err := repo.ResetStuckSessionsOnRestart(ctx) + require.NoError(t, err) + assert.Len(t, rows, 2) + for _, s := range rows { + assert.Equal(t, acp.SessionCrashed, s.Status) + } + rows2, _ := repo.ResetStuckSessionsOnRestart(ctx) + assert.Empty(t, rows2) +} + +func TestPgRepo_Event_OrderAndPurge(t *testing.T) { + t.Parallel() + ctx := context.Background() + repo, pool := setupRepo(t) + + uid := mustInsertUser(t, ctx, pool, "u5@local", false) + pid, wsid := mustInsertProjectWorkspace(t, ctx, pool, uid) + k, _ := repo.CreateAgentKind(ctx, &acp.AgentKind{ + ID: uuid.New(), Name: "kind5", DisplayName: "x", + BinaryPath: "x", Enabled: true, CreatedBy: uid, + }) + sess, _ := repo.InsertSession(ctx, &acp.Session{ + ID: uuid.New(), WorkspaceID: wsid, ProjectID: pid, + AgentKindID: k.ID, UserID: uid, + Branch: "main", CwdPath: "/tmp", IsMainWorktree: true, Status: acp.SessionRunning, + }) + + for i := 0; i < 3; i++ { + method := "session/update" + _, err := repo.InsertEvent(ctx, &acp.Event{ + SessionID: sess.ID, Direction: acp.DirectionOut, RPCKind: acp.RPCNotification, + Method: &method, Payload: []byte(`{"foo":"bar"}`), PayloadSize: 13, Truncated: false, + }) + require.NoError(t, err) + } + + // ListSince 0 → 全部 3 条 + all, err := repo.ListEventsSince(ctx, sess.ID, 0, 100) + require.NoError(t, err) + require.Len(t, all, 3) + assert.True(t, all[0].ID < all[1].ID && all[1].ID < all[2].ID) + + // Since 第 1 条 ID → 仅 2 条 + rest, _ := repo.ListEventsSince(ctx, sess.ID, all[0].ID, 100) + assert.Len(t, rest, 2) + + // Purge before now+1m → 全删 + n, _ := repo.PurgeEventsBefore(ctx, time.Now().Add(time.Minute)) + assert.Equal(t, int64(3), n) +} + +// ===== test helpers ===== + +func mustInsertUser(t *testing.T, ctx context.Context, pool *pgxpool.Pool, email string, admin bool) uuid.UUID { + t.Helper() + id := uuid.New() + _, err := pool.Exec(ctx, `INSERT INTO users (id, email, password_hash, display_name, is_admin) + VALUES ($1, $2, 'h', 'n', $3)`, id, email, admin) + require.NoError(t, err) + return id +} + +func mustInsertProjectWorkspace(t *testing.T, ctx context.Context, pool *pgxpool.Pool, owner uuid.UUID) (uuid.UUID, uuid.UUID) { + t.Helper() + pid := uuid.New() + wsid := uuid.New() + slug := "p-" + pid.String()[:8] + wsslug := "ws-" + wsid.String()[:8] + _, err := pool.Exec(ctx, `INSERT INTO projects (id, slug, name, owner_id) VALUES ($1, $2, 'p', $3)`, + pid, slug, owner) + require.NoError(t, err) + _, err = pool.Exec(ctx, `INSERT INTO workspaces (id, project_id, slug, name, git_remote_url, main_path) + VALUES ($1, $2, $3, 'w', 'git@x:y.git', '/tmp')`, wsid, pid, wsslug) + require.NoError(t, err) + return pid, wsid +}