238 lines
7.9 KiB
Go
238 lines
7.9 KiB
Go
package repository
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
|
|
"github.com/gochat/gochat/internal/model"
|
|
)
|
|
|
|
// createTestDocument is a helper that builds a CaptainDocument with typical fields.
|
|
func createTestDocument(accountID uint, assistantID uint, name string, externalLink string) *model.CaptainDocument {
|
|
metadata, _ := json.Marshal(map[string]interface{}{
|
|
"source": "web",
|
|
"version": 1,
|
|
})
|
|
|
|
return &model.CaptainDocument{
|
|
AccountID: accountID,
|
|
AssistantID: assistantID,
|
|
Name: name,
|
|
ExternalLink: externalLink,
|
|
Content: "This is test content for " + name,
|
|
Status: model.DocumentStatusInProgress,
|
|
SyncStatus: model.DocumentSyncStatusPending,
|
|
Metadata: metadata,
|
|
}
|
|
}
|
|
|
|
// --- 1. Create ---
|
|
func TestCaptainDocumentRepo_Create(t *testing.T) {
|
|
db := setupTestDB(t, &model.CaptainDocument{})
|
|
repo := NewCaptainDocumentRepo(db)
|
|
|
|
doc := createTestDocument(1, 10, "TestDoc", "https://example.com/doc1")
|
|
err := repo.Create(context.Background(), doc)
|
|
require.NoError(t, err)
|
|
assert.NotZero(t, doc.ID, "ID should be set after Create")
|
|
assert.Equal(t, "TestDoc", doc.Name)
|
|
assert.Equal(t, uint(1), doc.AccountID)
|
|
assert.Equal(t, uint(10), doc.AssistantID)
|
|
assert.Equal(t, model.DocumentStatusInProgress, doc.Status)
|
|
|
|
// Verify Metadata JSON was persisted
|
|
var foundMeta map[string]interface{}
|
|
require.NoError(t, json.Unmarshal(doc.Metadata, &foundMeta))
|
|
assert.Equal(t, "web", foundMeta["source"])
|
|
}
|
|
|
|
// --- 2. GetByID ---
|
|
func TestCaptainDocumentRepo_GetByID(t *testing.T) {
|
|
db := setupTestDB(t, &model.CaptainDocument{})
|
|
repo := NewCaptainDocumentRepo(db)
|
|
|
|
doc := createTestDocument(1, 10, "GetByIDDoc", "https://example.com/doc2")
|
|
err := repo.Create(context.Background(), doc)
|
|
require.NoError(t, err)
|
|
|
|
found, err := repo.GetByID(context.Background(), doc.ID)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, doc.ID, found.ID)
|
|
assert.Equal(t, "GetByIDDoc", found.Name)
|
|
assert.Equal(t, uint(1), found.AccountID)
|
|
assert.Equal(t, uint(10), found.AssistantID)
|
|
assert.Equal(t, "https://example.com/doc2", found.ExternalLink)
|
|
}
|
|
|
|
// --- 3. GetByID Not Found ---
|
|
func TestCaptainDocumentRepo_GetByID_NotFound(t *testing.T) {
|
|
db := setupTestDB(t, &model.CaptainDocument{})
|
|
repo := NewCaptainDocumentRepo(db)
|
|
|
|
found, err := repo.GetByID(context.Background(), 9999)
|
|
assert.Error(t, err)
|
|
assert.Nil(t, found)
|
|
}
|
|
|
|
// --- 4. Update ---
|
|
func TestCaptainDocumentRepo_Update(t *testing.T) {
|
|
db := setupTestDB(t, &model.CaptainDocument{})
|
|
repo := NewCaptainDocumentRepo(db)
|
|
|
|
doc := createTestDocument(1, 10, "BeforeUpdate", "https://example.com/doc3")
|
|
err := repo.Create(context.Background(), doc)
|
|
require.NoError(t, err)
|
|
|
|
doc.Name = "AfterUpdate"
|
|
doc.Content = "Updated content"
|
|
doc.Status = model.DocumentStatusCompleted
|
|
newMeta, _ := json.Marshal(map[string]interface{}{
|
|
"source": "api",
|
|
"version": 2,
|
|
})
|
|
doc.Metadata = newMeta
|
|
|
|
err = repo.Update(context.Background(), doc)
|
|
require.NoError(t, err)
|
|
|
|
found, err := repo.GetByID(context.Background(), doc.ID)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "AfterUpdate", found.Name)
|
|
assert.Equal(t, "Updated content", found.Content)
|
|
assert.Equal(t, model.DocumentStatusCompleted, found.Status)
|
|
|
|
// Verify Metadata JSON was updated
|
|
var foundMeta map[string]interface{}
|
|
require.NoError(t, json.Unmarshal(found.Metadata, &foundMeta))
|
|
assert.Equal(t, "api", foundMeta["source"])
|
|
assert.Equal(t, 2.0, foundMeta["version"])
|
|
}
|
|
|
|
// --- 5. Delete ---
|
|
func TestCaptainDocumentRepo_Delete(t *testing.T) {
|
|
db := setupTestDB(t, &model.CaptainDocument{})
|
|
repo := NewCaptainDocumentRepo(db)
|
|
|
|
doc := createTestDocument(1, 10, "DeleteMe", "https://example.com/doc4")
|
|
err := repo.Create(context.Background(), doc)
|
|
require.NoError(t, err)
|
|
|
|
err = repo.Delete(context.Background(), doc.ID)
|
|
require.NoError(t, err)
|
|
|
|
// After deletion, GetByID should return error
|
|
found, err := repo.GetByID(context.Background(), doc.ID)
|
|
assert.Error(t, err)
|
|
assert.Nil(t, found)
|
|
}
|
|
|
|
// --- 6. ListByAssistant ---
|
|
func TestCaptainDocumentRepo_ListByAssistant(t *testing.T) {
|
|
db := setupTestDB(t, &model.CaptainDocument{})
|
|
repo := NewCaptainDocumentRepo(db)
|
|
|
|
// Create 3 documents under assistant 10
|
|
for i := 0; i < 3; i++ {
|
|
d := createTestDocument(1, 10, "ListDoc"+string(rune('A'+i)), "https://example.com/list"+string(rune('A'+i)))
|
|
require.NoError(t, repo.Create(context.Background(), d))
|
|
}
|
|
// Create 1 document under assistant 20
|
|
d2 := createTestDocument(1, 20, "OtherAssistantDoc", "https://example.com/other")
|
|
require.NoError(t, repo.Create(context.Background(), d2))
|
|
|
|
docs, count, err := repo.ListByAssistant(context.Background(), 10, 0, 10)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, int64(3), count)
|
|
assert.Len(t, docs, 3)
|
|
|
|
for _, d := range docs {
|
|
assert.Equal(t, uint(10), d.AssistantID)
|
|
}
|
|
}
|
|
|
|
// --- 7. ListByAssistant with pagination ---
|
|
func TestCaptainDocumentRepo_ListByAssistant_Pagination(t *testing.T) {
|
|
db := setupTestDB(t, &model.CaptainDocument{})
|
|
repo := NewCaptainDocumentRepo(db)
|
|
|
|
// Create 5 documents under assistant 10
|
|
for i := 0; i < 5; i++ {
|
|
d := createTestDocument(1, 10, "PagDoc"+string(rune('A'+i)), "https://example.com/pag"+string(rune('A'+i)))
|
|
require.NoError(t, repo.Create(context.Background(), d))
|
|
}
|
|
|
|
// Fetch first page: offset=0, limit=2
|
|
docs, count, err := repo.ListByAssistant(context.Background(), 10, 0, 2)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, int64(5), count, "total count should be 5 regardless of pagination")
|
|
assert.Len(t, docs, 2, "first page should return 2 items")
|
|
|
|
// Fetch second page: offset=2, limit=2
|
|
docs2, count2, err := repo.ListByAssistant(context.Background(), 10, 2, 2)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, int64(5), count2)
|
|
assert.Len(t, docs2, 2, "second page should return 2 items")
|
|
}
|
|
|
|
// --- 8. FindByExternalLink ---
|
|
func TestCaptainDocumentRepo_FindByExternalLink(t *testing.T) {
|
|
db := setupTestDB(t, &model.CaptainDocument{})
|
|
repo := NewCaptainDocumentRepo(db)
|
|
|
|
doc := createTestDocument(1, 10, "ExternalLinkDoc", "https://example.com/unique-link")
|
|
err := repo.Create(context.Background(), doc)
|
|
require.NoError(t, err)
|
|
|
|
found, err := repo.FindByExternalLink(context.Background(), 10, "https://example.com/unique-link")
|
|
require.NoError(t, err)
|
|
assert.Equal(t, doc.ID, found.ID)
|
|
assert.Equal(t, "ExternalLinkDoc", found.Name)
|
|
assert.Equal(t, "https://example.com/unique-link", found.ExternalLink)
|
|
}
|
|
|
|
// --- 9. FindByExternalLink Not Found ---
|
|
func TestCaptainDocumentRepo_FindByExternalLink_NotFound(t *testing.T) {
|
|
db := setupTestDB(t, &model.CaptainDocument{})
|
|
repo := NewCaptainDocumentRepo(db)
|
|
|
|
// Create a document with a different external link
|
|
doc := createTestDocument(1, 10, "SomeDoc", "https://example.com/existing-link")
|
|
err := repo.Create(context.Background(), doc)
|
|
require.NoError(t, err)
|
|
|
|
// Search for a non-existent external link
|
|
found, err := repo.FindByExternalLink(context.Background(), 10, "https://example.com/nonexistent")
|
|
assert.Error(t, err)
|
|
assert.Nil(t, found)
|
|
}
|
|
|
|
// --- 10. FindByStatus ---
|
|
func TestCaptainDocumentRepo_FindByStatus(t *testing.T) {
|
|
db := setupTestDB(t, &model.CaptainDocument{})
|
|
repo := NewCaptainDocumentRepo(db)
|
|
|
|
// Create 2 in_progress documents
|
|
for i := 0; i < 2; i++ {
|
|
d := createTestDocument(1, 10, "InProgressDoc"+string(rune('A'+i)), "https://example.com/ip"+string(rune('A'+i)))
|
|
d.Status = model.DocumentStatusInProgress
|
|
require.NoError(t, repo.Create(context.Background(), d))
|
|
}
|
|
// Create 1 completed document
|
|
d3 := createTestDocument(1, 10, "CompletedDoc", "https://example.com/completed")
|
|
d3.Status = model.DocumentStatusCompleted
|
|
require.NoError(t, repo.Create(context.Background(), d3))
|
|
|
|
docs, count, err := repo.FindByStatus(context.Background(), 10, model.DocumentStatusInProgress, 0, 10)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, int64(2), count)
|
|
assert.Len(t, docs, 2)
|
|
|
|
for _, d := range docs {
|
|
assert.Equal(t, model.DocumentStatusInProgress, d.Status)
|
|
assert.Equal(t, uint(10), d.AssistantID)
|
|
}
|
|
} |