Files
gochat/internal/repository/copilot_thread_repo_test.go
T
2026-06-04 15:44:48 +08:00

196 lines
6.4 KiB
Go

package repository
import (
"context"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/gochat/gochat/internal/model"
)
// --- 1. Create ---
func TestCopilotThreadRepo_Create(t *testing.T) {
db := setupTestDB(t, &model.CopilotThread{}, &model.CopilotMessage{})
repo := NewCopilotThreadRepo(db)
thread := &model.CopilotThread{
AccountID: 1,
UserID: 10,
Title: "My Copilot Thread",
}
err := repo.Create(context.Background(), thread)
require.NoError(t, err)
assert.NotZero(t, thread.ID, "ID should be set after Create")
assert.Equal(t, uint(1), thread.AccountID)
assert.Equal(t, uint(10), thread.UserID)
assert.Equal(t, "My Copilot Thread", thread.Title)
}
// --- 2. GetByID ---
func TestCopilotThreadRepo_GetByID(t *testing.T) {
db := setupTestDB(t, &model.CopilotThread{}, &model.CopilotMessage{})
repo := NewCopilotThreadRepo(db)
thread := createTestCopilotThread(t, db, 1, 10, "Test Thread")
found, err := repo.GetByID(context.Background(), thread.ID)
require.NoError(t, err)
assert.Equal(t, thread.ID, found.ID)
assert.Equal(t, thread.AccountID, found.AccountID)
assert.Equal(t, thread.UserID, found.UserID)
assert.Equal(t, "Test Thread", found.Title)
}
func TestCopilotThreadRepo_GetByID_NotFound(t *testing.T) {
db := setupTestDB(t, &model.CopilotThread{}, &model.CopilotMessage{})
repo := NewCopilotThreadRepo(db)
found, err := repo.GetByID(context.Background(), 9999)
assert.Error(t, err)
assert.Nil(t, found)
}
// --- 3. Update ---
func TestCopilotThreadRepo_Update(t *testing.T) {
db := setupTestDB(t, &model.CopilotThread{}, &model.CopilotMessage{})
repo := NewCopilotThreadRepo(db)
thread := createTestCopilotThread(t, db, 1, 10, "Original Title")
thread.Title = "Updated Title"
err := repo.Update(context.Background(), thread)
require.NoError(t, err)
found, err := repo.GetByID(context.Background(), thread.ID)
require.NoError(t, err)
assert.Equal(t, "Updated Title", found.Title)
}
// --- 4. Delete ---
func TestCopilotThreadRepo_Delete(t *testing.T) {
db := setupTestDB(t, &model.CopilotThread{}, &model.CopilotMessage{})
repo := NewCopilotThreadRepo(db)
thread := createTestCopilotThread(t, db, 1, 10, "To Be Deleted")
err := repo.Delete(context.Background(), thread.ID)
require.NoError(t, err)
found, err := repo.GetByID(context.Background(), thread.ID)
assert.Error(t, err)
assert.Nil(t, found)
}
// --- 5. ListByUser ---
func TestCopilotThreadRepo_ListByUser(t *testing.T) {
db := setupTestDB(t, &model.CopilotThread{}, &model.CopilotMessage{})
repo := NewCopilotThreadRepo(db)
// Create threads for user 10 under account 1
createTestCopilotThread(t, db, 1, 10, "Thread A")
createTestCopilotThread(t, db, 1, 10, "Thread B")
// Create a thread for a different user — should NOT appear
createTestCopilotThread(t, db, 1, 20, "Other User Thread")
threads, count, err := repo.ListByUser(context.Background(), 1, 10, 0, 10)
require.NoError(t, err)
assert.Equal(t, int64(2), count)
assert.Len(t, threads, 2)
// Verify pagination: offset=1, limit=1 should return exactly 1 thread
threads2, count2, err := repo.ListByUser(context.Background(), 1, 10, 1, 1)
require.NoError(t, err)
assert.Equal(t, int64(2), count2)
assert.Len(t, threads2, 1)
}
// --- 6. FindByAccountID ---
func TestCopilotThreadRepo_FindByAccountID(t *testing.T) {
db := setupTestDB(t, &model.CopilotThread{}, &model.CopilotMessage{})
repo := NewCopilotThreadRepo(db)
// Create threads across two accounts
createTestCopilotThread(t, db, 1, 10, "Account1 Thread1")
createTestCopilotThread(t, db, 1, 20, "Account1 Thread2")
createTestCopilotThread(t, db, 2, 10, "Account2 Thread1")
threads, count, err := repo.FindByAccountID(context.Background(), 1, 0, 10)
require.NoError(t, err)
assert.Equal(t, int64(2), count)
assert.Len(t, threads, 2)
for _, th := range threads {
assert.Equal(t, uint(1), th.AccountID)
}
// Account 2 should have exactly 1 thread
threads2, count2, err := repo.FindByAccountID(context.Background(), 2, 0, 10)
require.NoError(t, err)
assert.Equal(t, int64(1), count2)
assert.Len(t, threads2, 1)
}
// --- 7. FindByAssistantID ---
func TestCopilotThreadRepo_FindByAssistantID(t *testing.T) {
db := setupTestDB(t, &model.CopilotThread{}, &model.CopilotMessage{})
repo := NewCopilotThreadRepo(db)
assistantID1 := uint(100)
assistantID2 := uint(200)
// Create threads linked to assistant 100
th1 := &model.CopilotThread{AccountID: 1, UserID: 10, AssistantID: &assistantID1, Title: "With Assistant 100"}
require.NoError(t, db.Create(th1).Error)
th2 := &model.CopilotThread{AccountID: 1, UserID: 20, AssistantID: &assistantID1, Title: "With Assistant 100 (2)"}
require.NoError(t, db.Create(th2).Error)
// Create a thread linked to assistant 200
th3 := &model.CopilotThread{AccountID: 2, UserID: 30, AssistantID: &assistantID2, Title: "With Assistant 200"}
require.NoError(t, db.Create(th3).Error)
// Create a thread with no assistant
createTestCopilotThread(t, db, 1, 40, "No Assistant")
// Find threads for assistant 100
threads, count, err := repo.FindByAssistantID(context.Background(), assistantID1, 0, 10)
require.NoError(t, err)
assert.Equal(t, int64(2), count)
assert.Len(t, threads, 2)
for _, th := range threads {
assert.NotNil(t, th.AssistantID)
assert.Equal(t, assistantID1, *th.AssistantID)
}
// Find threads for assistant 200
threads2, count2, err := repo.FindByAssistantID(context.Background(), assistantID2, 0, 10)
require.NoError(t, err)
assert.Equal(t, int64(1), count2)
assert.Len(t, threads2, 1)
assert.Equal(t, assistantID2, *threads2[0].AssistantID)
}
// --- 8. GetByID with Messages (Preload) ---
func TestCopilotThreadRepo_GetByID_PreloadsMessages(t *testing.T) {
db := setupTestDB(t, &model.CopilotThread{}, &model.CopilotMessage{})
msgRepo := NewCopilotMessageRepo(db)
threadRepo := NewCopilotThreadRepo(db)
thread := createTestCopilotThread(t, db, 1, 10, "Thread With Messages")
// Add messages to the thread
msg1 := createTestCopilotMessage(1, thread.ID, model.CopilotMessageTypeUser, "Hello")
require.NoError(t, msgRepo.Create(context.Background(), msg1))
msg2 := createTestCopilotMessage(1, thread.ID, model.CopilotMessageTypeAssistant, "Hi there")
require.NoError(t, msgRepo.Create(context.Background(), msg2))
found, err := threadRepo.GetByID(context.Background(), thread.ID)
require.NoError(t, err)
assert.Equal(t, thread.ID, found.ID)
assert.Len(t, found.Messages, 2, "GetByID should preload Messages")
}