agentic-coding-workflow/internal/chat/repository_test.go

//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)
}