//go:build integration package chat import ( "context" "strings" "testing" "time" "github.com/google/uuid" "github.com/jackc/pgx/v5/pgxpool" "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" ) // setupRepo 起一个临时 postgres 容器、跑迁移、注入 user / llm_endpoint / llm_model // 行作外键基础,返回 Repository 与 pool、uid、eid、mid。 func setupRepo(t *testing.T) (Repository, *pgxpool.Pool, uuid.UUID, uuid.UUID, uuid.UUID) { 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) uid := uuid.New() _, err = pool.Exec(ctx, `INSERT INTO users (id,email,password_hash,display_name,is_admin,created_at,updated_at) VALUES ($1,'a@b.c','x','admin',true,now(),now())`, uid) require.NoError(t, err) eid := uuid.New() _, err = pool.Exec(ctx, `INSERT INTO llm_endpoints (id,provider,display_name,base_url,api_key_encrypted,created_by) VALUES ($1,'anthropic','test','https://example.com','\x',$2)`, eid, uid) require.NoError(t, err) mid := uuid.New() _, err = pool.Exec(ctx, `INSERT INTO llm_models (id,endpoint_id,model_id,display_name,capabilities,context_window,max_output_tokens) VALUES ($1,$2,'claude','Claude','{}',200000,8192)`, mid, eid) require.NoError(t, err) return NewRepository(pool), pool, uid, eid, mid } // ===== 1. Conversation CRUD ===== func TestRepo_Conversation_CRUD(t *testing.T) { t.Parallel() repo, _, uid, _, mid := setupRepo(t) ctx := context.Background() cid := uuid.New() conv, err := repo.InsertConversation(ctx, InsertConversationParams{ ID: cid, UserID: uid, ModelID: mid, SystemPrompt: "be helpful", }) require.NoError(t, err) require.Equal(t, cid, conv.ID) require.Equal(t, "be helpful", conv.SystemPrompt) got, err := repo.GetConversation(ctx, cid) require.NoError(t, err) require.Equal(t, cid, got.ID) require.NoError(t, repo.SoftDeleteConversation(ctx, cid)) _, err = repo.GetConversation(ctx, cid) var ae *errs.AppError require.ErrorAs(t, err, &ae) require.Equal(t, errs.CodeChatConversationNotFound, ae.Code) } // ===== 2. Messages: pending unique index ===== func TestRepo_Messages_PendingUniqueIndex(t *testing.T) { t.Parallel() repo, _, uid, _, mid := setupRepo(t) ctx := context.Background() cid := uuid.New() _, err := repo.InsertConversation(ctx, InsertConversationParams{ ID: cid, UserID: uid, ModelID: mid, SystemPrompt: "", }) require.NoError(t, err) // first pending assistant message — should succeed _, err = repo.InsertMessage(ctx, cid, RoleAssistant, "", StatusPending) require.NoError(t, err) // second pending assistant message in same conv — must fail with CodeChatConcurrentMessage _, err = repo.InsertMessage(ctx, cid, RoleAssistant, "", StatusPending) require.Error(t, err) var ae *errs.AppError require.ErrorAs(t, err, &ae) require.Equal(t, errs.CodeChatConcurrentMessage, ae.Code) require.True(t, strings.Contains(ae.Message, "another message is already streaming")) } // ===== 3. Message status transitions ===== func TestRepo_Messages_StatusTransitions(t *testing.T) { t.Parallel() repo, _, uid, _, mid := setupRepo(t) ctx := context.Background() cid := uuid.New() _, err := repo.InsertConversation(ctx, InsertConversationParams{ ID: cid, UserID: uid, ModelID: mid, SystemPrompt: "", }) require.NoError(t, err) msg, err := repo.InsertMessage(ctx, cid, RoleAssistant, "", StatusPending) require.NoError(t, err) require.Equal(t, StatusPending, msg.Status) pt := int32(10) ct := int32(20) require.NoError(t, repo.UpdateMessageOK(ctx, UpdateMessageOKParams{ ID: msg.ID, Content: "hello", PromptTokens: &pt, CompletionTokens: &ct, })) updated, err := repo.GetMessage(ctx, msg.ID) require.NoError(t, err) require.Equal(t, StatusOK, updated.Status) require.Equal(t, "hello", updated.Content) require.NotNil(t, updated.PromptTokens) require.Equal(t, 10, *updated.PromptTokens) } // ===== 4. Attachments orphan lifecycle ===== func TestRepo_Attachments_OrphanLifecycle(t *testing.T) { t.Parallel() repo, _, uid, _, mid := setupRepo(t) ctx := context.Background() aid := uuid.New() att, err := repo.InsertAttachment(ctx, InsertAttachmentParams{ ID: aid, UserID: uid, Filename: "file.txt", MimeType: "text/plain", SizeBytes: 42, SHA256: "abc", StoragePath: "/tmp/file", }) require.NoError(t, err) require.Nil(t, att.MessageID, "newly inserted attachment must be orphan") // create conversation + message to attach to cid := uuid.New() _, err = repo.InsertConversation(ctx, InsertConversationParams{ ID: cid, UserID: uid, ModelID: mid, SystemPrompt: "", }) require.NoError(t, err) msg, err := repo.InsertMessage(ctx, cid, RoleUser, "hello", StatusOK) require.NoError(t, err) require.NoError(t, repo.AttachToMessage(ctx, []uuid.UUID{aid}, msg.ID)) atts, err := repo.ListAttachmentsForMessage(ctx, msg.ID) require.NoError(t, err) require.Len(t, atts, 1) require.Equal(t, aid, atts[0].ID) require.NotNil(t, atts[0].MessageID) require.Equal(t, msg.ID, *atts[0].MessageID) } // ===== 5. Templates scope filtering ===== func TestRepo_Templates_ScopeFiltering(t *testing.T) { t.Parallel() repo, pool, uid, _, _ := setupRepo(t) ctx := context.Background() // seed user B uidB := uuid.New() _, err := pool.Exec(ctx, `INSERT INTO users (id,email,password_hash,display_name,is_admin,created_at,updated_at) VALUES ($1,'b@b.c','x','B',false,now(),now())`, uidB) require.NoError(t, err) // system template (no owner) _, err = repo.InsertPromptTemplate(ctx, InsertTemplateParams{ ID: uuid.New(), Name: "sys", Content: "system", Scope: TemplateScopeSystem, }) require.NoError(t, err) // user A template _, err = repo.InsertPromptTemplate(ctx, InsertTemplateParams{ ID: uuid.New(), Name: "tmpl-a", Content: "for a", Scope: TemplateScopeUser, OwnerID: &uid, }) require.NoError(t, err) // user B template _, err = repo.InsertPromptTemplate(ctx, InsertTemplateParams{ ID: uuid.New(), Name: "tmpl-b", Content: "for b", Scope: TemplateScopeUser, OwnerID: &uidB, }) require.NoError(t, err) templates, err := repo.ListPromptTemplatesForUser(ctx, uid) require.NoError(t, err) require.Len(t, templates, 2, "user A should see system + own template only") names := make([]string, len(templates)) for i, tmpl := range templates { names[i] = tmpl.Name } require.Contains(t, names, "sys") require.Contains(t, names, "tmpl-a") require.NotContains(t, names, "tmpl-b") } // ===== 6. Endpoints CRUD ===== func TestRepo_Endpoints_CRUD(t *testing.T) { t.Parallel() repo, _, uid, _, _ := setupRepo(t) ctx := context.Background() eid := uuid.New() ep, err := repo.InsertLLMEndpoint(ctx, InsertEndpointParams{ ID: eid, Provider: ProviderOpenAI, DisplayName: "OpenAI", BaseURL: "https://api.openai.com", APIKeyEncrypted: []byte("key"), CreatedBy: uid, }) require.NoError(t, err) require.Equal(t, eid, ep.ID) require.Equal(t, ProviderOpenAI, ep.Provider) newName := "OpenAI Updated" require.NoError(t, repo.UpdateLLMEndpoint(ctx, eid, UpdateEndpointInput{ DisplayName: &newName, })) got, err := repo.GetLLMEndpoint(ctx, eid) require.NoError(t, err) require.Equal(t, "OpenAI Updated", got.DisplayName) require.NoError(t, repo.DeleteLLMEndpoint(ctx, eid)) _, err = repo.GetLLMEndpoint(ctx, eid) var ae *errs.AppError require.ErrorAs(t, err, &ae) require.Equal(t, errs.CodeChatEndpointNotFound, ae.Code) } // ===== 7. Models: title generator unique constraint ===== func TestRepo_Models_TitleGeneratorUnique(t *testing.T) { t.Parallel() repo, _, _, eid, _ := setupRepo(t) ctx := context.Background() // first title-generator model _, err := repo.InsertLLMModel(ctx, InsertModelParams{ ID: uuid.New(), EndpointID: eid, ModelID: "m1", DisplayName: "M1", Capabilities: ModelCapabilities{}, ContextWindow: 4096, MaxOutputTokens: 1024, IsTitleGenerator: true, Enabled: true, }) require.NoError(t, err) // second title-generator model — must fail with CodeConflict _, err = repo.InsertLLMModel(ctx, InsertModelParams{ ID: uuid.New(), EndpointID: eid, ModelID: "m2", DisplayName: "M2", Capabilities: ModelCapabilities{}, ContextWindow: 4096, MaxOutputTokens: 1024, IsTitleGenerator: true, Enabled: true, }) require.Error(t, err) var ae *errs.AppError require.ErrorAs(t, err, &ae) require.Equal(t, errs.CodeConflict, ae.Code) } // ===== 8. Usage summarize by day ===== func TestRepo_Usage_SummarizeByDay(t *testing.T) { t.Parallel() repo, _, uid, eid, mid := setupRepo(t) ctx := context.Background() cid := uuid.New() _, err := repo.InsertConversation(ctx, InsertConversationParams{ ID: cid, UserID: uid, ModelID: mid, SystemPrompt: "", }) require.NoError(t, err) msg, err := repo.InsertMessage(ctx, cid, RoleUser, "hi", StatusOK) require.NoError(t, err) // insert 3 usage records for i := 0; i < 3; i++ { require.NoError(t, repo.InsertUsage(ctx, UsageRecord{ ConversationID: cid, MessageID: msg.ID, UserID: uid, EndpointID: eid, ModelID: mid, PromptTokens: 10, CompletionTokens: 20, ThinkingTokens: 0, CostUSD: 0.001, })) } from := time.Now().UTC().Truncate(24 * time.Hour) to := from.Add(24 * time.Hour) rows, err := repo.SummarizeUsageByDay(ctx, from, to) require.NoError(t, err) require.Len(t, rows, 1) require.Equal(t, 30, rows[0].PromptTokens) require.Equal(t, 60, rows[0].CompletionTokens) } // ===== 9. Tx rollback undoes insert ===== func TestRepo_Tx_RollbackUndoesInsert(t *testing.T) { t.Parallel() repo, _, uid, _, mid := setupRepo(t) ctx := context.Background() txRepo, tx, err := repo.WithTx(ctx) require.NoError(t, err) cid := uuid.New() _, err = txRepo.InsertConversation(ctx, InsertConversationParams{ ID: cid, UserID: uid, ModelID: mid, SystemPrompt: "rollback me", }) require.NoError(t, err) // can read within same tx got, err := txRepo.GetConversation(ctx, cid) require.NoError(t, err) require.Equal(t, cid, got.ID) // rollback require.NoError(t, tx.Rollback(ctx)) // after rollback, reading via original repo should return not found _, err = repo.GetConversation(ctx, cid) var ae *errs.AppError require.ErrorAs(t, err, &ae) require.Equal(t, errs.CodeChatConversationNotFound, ae.Code) }