01820d9fe2
- Introduced `AddFunc` method to the `Job` struct for adding Go functions as job executions, allowing for dynamic execution of functions with specified arguments. - Enhanced internal execution handling to register functions in a global registry, ensuring proper cleanup after execution. - Implemented `ExecuteFunc` method in the `Goroutine` struct to handle the execution of registered functions, including error handling and context management. - Added comprehensive unit tests for `AddFunc`, verifying function registration, execution, and memory cleanup post-execution. - Updated related documentation to reflect the new functionality and usage patterns for adding and executing Go functions within the job system.
485 lines
14 KiB
Go
485 lines
14 KiB
Go
package api_test
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
graphragtypes "github.com/yaoapp/gou/graphrag/types"
|
|
"github.com/yaoapp/yao/kb"
|
|
"github.com/yaoapp/yao/kb/api"
|
|
)
|
|
|
|
// Note: TestMain is defined in collection_test.go, which handles environment setup
|
|
// Run tests with: source env.local.sh && go test -v ./kb/api/...
|
|
|
|
// createTestCollectionForText is a helper to create a test collection for text tests
|
|
func createTestCollectionForText(t *testing.T, ctx context.Context) string {
|
|
if kb.API == nil {
|
|
t.Skip("KB API not initialized")
|
|
}
|
|
|
|
collectionID := fmt.Sprintf("test_text_%d", time.Now().UnixNano())
|
|
|
|
params := &api.CreateCollectionParams{
|
|
ID: collectionID,
|
|
Metadata: map[string]interface{}{
|
|
"name": "Test Text Collection",
|
|
"description": "Collection for AddText tests",
|
|
},
|
|
EmbeddingProviderID: "__yao.openai",
|
|
EmbeddingOptionID: "text-embedding-3-small",
|
|
Locale: "en",
|
|
Config: &graphragtypes.CreateCollectionOptions{
|
|
Distance: "cosine",
|
|
IndexType: "hnsw",
|
|
},
|
|
}
|
|
|
|
_, err := kb.API.CreateCollection(ctx, params)
|
|
if err != nil {
|
|
t.Fatalf("Failed to create test collection: %v", err)
|
|
}
|
|
|
|
return collectionID
|
|
}
|
|
|
|
// cleanupTestCollectionForText removes a test collection
|
|
func cleanupTestCollectionForText(ctx context.Context, collectionID string) {
|
|
if kb.API != nil {
|
|
_, _ = kb.API.RemoveCollection(ctx, collectionID)
|
|
}
|
|
}
|
|
|
|
// ========== AddText Tests ==========
|
|
|
|
func TestAddText(t *testing.T) {
|
|
if kb.API == nil {
|
|
t.Skip("KB API not initialized")
|
|
}
|
|
|
|
ctx := context.Background()
|
|
collectionID := createTestCollectionForText(t, ctx)
|
|
defer cleanupTestCollectionForText(ctx, collectionID)
|
|
|
|
t.Run("AddTextSuccess", func(t *testing.T) {
|
|
params := &api.AddTextParams{
|
|
CollectionID: collectionID,
|
|
Text: "This is a test document content for knowledge base testing. It contains some sample text that will be chunked and embedded.",
|
|
Locale: "en",
|
|
Metadata: map[string]interface{}{
|
|
"title": "Test Text Document",
|
|
"description": "A test document",
|
|
},
|
|
Chunking: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.structured",
|
|
OptionID: "standard",
|
|
},
|
|
Embedding: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.openai",
|
|
OptionID: "text-embedding-3-small",
|
|
},
|
|
AuthScope: map[string]interface{}{
|
|
"__yao_created_by": "test_user",
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.AddText(ctx, params)
|
|
if err != nil {
|
|
// Skip if connector not loaded (environment issue)
|
|
if assert.Contains(t, err.Error(), "connector") {
|
|
t.Skipf("Skipping due to connector not loaded: %v", err)
|
|
}
|
|
}
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, result)
|
|
if result != nil {
|
|
assert.Equal(t, collectionID, result.CollectionID)
|
|
assert.NotEmpty(t, result.DocID)
|
|
assert.Contains(t, result.Message, "successfully")
|
|
t.Logf("Added text document: %s", result.DocID)
|
|
|
|
// Verify document was created
|
|
doc, err := kb.API.GetDocument(ctx, result.DocID, nil)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, doc)
|
|
assert.Equal(t, "text", doc["type"])
|
|
assert.Equal(t, "Test Text Document", doc["name"])
|
|
assert.Equal(t, "completed", doc["status"])
|
|
t.Logf("✅ Document verified: type=%v, name=%v, status=%v", doc["type"], doc["name"], doc["status"])
|
|
}
|
|
})
|
|
|
|
t.Run("AddTextMissingCollectionID", func(t *testing.T) {
|
|
params := &api.AddTextParams{
|
|
Text: "Some text content",
|
|
Chunking: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.structured",
|
|
},
|
|
Embedding: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.openai",
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.AddText(ctx, params)
|
|
assert.Error(t, err)
|
|
assert.Nil(t, result)
|
|
assert.Contains(t, err.Error(), "collection_id is required")
|
|
})
|
|
|
|
t.Run("AddTextMissingText", func(t *testing.T) {
|
|
params := &api.AddTextParams{
|
|
CollectionID: collectionID,
|
|
Chunking: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.structured",
|
|
},
|
|
Embedding: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.openai",
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.AddText(ctx, params)
|
|
assert.Error(t, err)
|
|
assert.Nil(t, result)
|
|
assert.Contains(t, err.Error(), "text is required")
|
|
})
|
|
|
|
t.Run("AddTextMissingChunking", func(t *testing.T) {
|
|
params := &api.AddTextParams{
|
|
CollectionID: collectionID,
|
|
Text: "Some text content",
|
|
Embedding: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.openai",
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.AddText(ctx, params)
|
|
assert.Error(t, err)
|
|
assert.Nil(t, result)
|
|
assert.Contains(t, err.Error(), "chunking configuration is required")
|
|
})
|
|
|
|
t.Run("AddTextMissingEmbedding", func(t *testing.T) {
|
|
params := &api.AddTextParams{
|
|
CollectionID: collectionID,
|
|
Text: "Some text content",
|
|
Chunking: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.structured",
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.AddText(ctx, params)
|
|
assert.Error(t, err)
|
|
assert.Nil(t, result)
|
|
assert.Contains(t, err.Error(), "embedding configuration is required")
|
|
})
|
|
|
|
t.Run("AddTextWithCustomDocID", func(t *testing.T) {
|
|
customDocID := fmt.Sprintf("custom_text_doc_%d", time.Now().UnixNano())
|
|
params := &api.AddTextParams{
|
|
CollectionID: collectionID,
|
|
DocID: customDocID,
|
|
Text: "Text with custom document ID",
|
|
Chunking: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.structured",
|
|
OptionID: "standard",
|
|
},
|
|
Embedding: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.openai",
|
|
OptionID: "text-embedding-3-small",
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.AddText(ctx, params)
|
|
if err != nil {
|
|
t.Skipf("Skipping due to error: %v", err)
|
|
}
|
|
assert.NotNil(t, result)
|
|
if result != nil {
|
|
assert.Equal(t, customDocID, result.DocID)
|
|
t.Logf("Added text with custom DocID: %s", result.DocID)
|
|
}
|
|
})
|
|
|
|
t.Run("AddTextWithTitleFromMetadata", func(t *testing.T) {
|
|
params := &api.AddTextParams{
|
|
CollectionID: collectionID,
|
|
Text: "Text content with title from metadata",
|
|
Metadata: map[string]interface{}{
|
|
"title": "Custom Title From Metadata",
|
|
},
|
|
Chunking: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.structured",
|
|
OptionID: "standard",
|
|
},
|
|
Embedding: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.openai",
|
|
OptionID: "text-embedding-3-small",
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.AddText(ctx, params)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, result)
|
|
|
|
if result != nil {
|
|
doc, err := kb.API.GetDocument(ctx, result.DocID, nil)
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, "Custom Title From Metadata", doc["name"])
|
|
t.Logf("✅ Title from metadata verified: %v", doc["name"])
|
|
}
|
|
})
|
|
}
|
|
|
|
// ========== AddTextAsync Tests ==========
|
|
|
|
func TestAddTextAsync(t *testing.T) {
|
|
if kb.API == nil {
|
|
t.Skip("KB API not initialized")
|
|
}
|
|
|
|
ctx := context.Background()
|
|
collectionID := createTestCollectionForText(t, ctx)
|
|
defer cleanupTestCollectionForText(ctx, collectionID)
|
|
|
|
t.Run("AddTextAsyncSuccess", func(t *testing.T) {
|
|
params := &api.AddTextParams{
|
|
CollectionID: collectionID,
|
|
Text: "This is async text content for testing background processing.",
|
|
Locale: "en",
|
|
Metadata: map[string]interface{}{
|
|
"title": "Async Text Document",
|
|
},
|
|
Chunking: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.structured",
|
|
OptionID: "standard",
|
|
},
|
|
Embedding: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.openai",
|
|
OptionID: "text-embedding-3-small",
|
|
},
|
|
Job: &api.JobOptionsParams{
|
|
Name: "Test Async Text Job",
|
|
Description: "Testing async text processing",
|
|
Category: "Test",
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.AddTextAsync(ctx, params)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, result)
|
|
if result != nil {
|
|
assert.NotEmpty(t, result.JobID)
|
|
assert.NotEmpty(t, result.DocID)
|
|
t.Logf("Created async job: %s for document: %s", result.JobID, result.DocID)
|
|
}
|
|
|
|
// Verify document was created with pending status
|
|
if result != nil {
|
|
doc, err := kb.API.GetDocument(ctx, result.DocID, nil)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, doc)
|
|
assert.Equal(t, "text", doc["type"])
|
|
assert.Equal(t, result.JobID, doc["job_id"])
|
|
t.Logf("✅ Async document created: status=%v, job_id=%v", doc["status"], doc["job_id"])
|
|
|
|
// Wait for job to complete (max 30 seconds)
|
|
maxWait := 30 * time.Second
|
|
pollInterval := 500 * time.Millisecond
|
|
startTime := time.Now()
|
|
var finalStatus string
|
|
|
|
for time.Since(startTime) < maxWait {
|
|
doc, err = kb.API.GetDocument(ctx, result.DocID, nil)
|
|
if err != nil {
|
|
t.Logf("Error getting document: %v", err)
|
|
break
|
|
}
|
|
finalStatus, _ = doc["status"].(string)
|
|
if finalStatus == "completed" || finalStatus == "error" {
|
|
break
|
|
}
|
|
time.Sleep(pollInterval)
|
|
}
|
|
|
|
t.Logf("✅ Job completed: final status=%s, elapsed=%v", finalStatus, time.Since(startTime))
|
|
if finalStatus == "error" {
|
|
if errMsg, ok := doc["error_message"].(string); ok {
|
|
t.Logf("Error message: %s", errMsg)
|
|
}
|
|
}
|
|
assert.Equal(t, "completed", finalStatus, "Job should complete successfully")
|
|
}
|
|
})
|
|
|
|
t.Run("AddTextAsyncMissingCollectionID", func(t *testing.T) {
|
|
params := &api.AddTextParams{
|
|
Text: "Some text",
|
|
Chunking: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.structured",
|
|
},
|
|
Embedding: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.openai",
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.AddTextAsync(ctx, params)
|
|
assert.Error(t, err)
|
|
assert.Nil(t, result)
|
|
assert.Contains(t, err.Error(), "collection_id is required")
|
|
})
|
|
|
|
t.Run("AddTextAsyncMissingText", func(t *testing.T) {
|
|
params := &api.AddTextParams{
|
|
CollectionID: collectionID,
|
|
Chunking: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.structured",
|
|
},
|
|
Embedding: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.openai",
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.AddTextAsync(ctx, params)
|
|
assert.Error(t, err)
|
|
assert.Nil(t, result)
|
|
assert.Contains(t, err.Error(), "text is required")
|
|
})
|
|
|
|
t.Run("AddTextAsyncMissingChunking", func(t *testing.T) {
|
|
params := &api.AddTextParams{
|
|
CollectionID: collectionID,
|
|
Text: "Some text",
|
|
Embedding: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.openai",
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.AddTextAsync(ctx, params)
|
|
assert.Error(t, err)
|
|
assert.Nil(t, result)
|
|
assert.Contains(t, err.Error(), "chunking configuration is required")
|
|
})
|
|
|
|
t.Run("AddTextAsyncMissingEmbedding", func(t *testing.T) {
|
|
params := &api.AddTextParams{
|
|
CollectionID: collectionID,
|
|
Text: "Some text",
|
|
Chunking: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.structured",
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.AddTextAsync(ctx, params)
|
|
assert.Error(t, err)
|
|
assert.Nil(t, result)
|
|
assert.Contains(t, err.Error(), "embedding configuration is required")
|
|
})
|
|
|
|
t.Run("AddTextAsyncWithCustomDocID", func(t *testing.T) {
|
|
customDocID := fmt.Sprintf("async_custom_text_%d", time.Now().UnixNano())
|
|
params := &api.AddTextParams{
|
|
CollectionID: collectionID,
|
|
DocID: customDocID,
|
|
Text: "Async text with custom DocID",
|
|
Chunking: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.structured",
|
|
OptionID: "standard",
|
|
},
|
|
Embedding: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.openai",
|
|
OptionID: "text-embedding-3-small",
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.AddTextAsync(ctx, params)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, result)
|
|
if result != nil {
|
|
assert.Equal(t, customDocID, result.DocID)
|
|
t.Logf("Created async text with custom DocID: %s", result.DocID)
|
|
}
|
|
})
|
|
}
|
|
|
|
// ========== AddText Integration Test ==========
|
|
|
|
func TestAddTextIntegration(t *testing.T) {
|
|
if kb.API == nil {
|
|
t.Skip("KB API not initialized")
|
|
}
|
|
|
|
ctx := context.Background()
|
|
collectionID := createTestCollectionForText(t, ctx)
|
|
defer cleanupTestCollectionForText(ctx, collectionID)
|
|
|
|
t.Run("FullTextLifecycle", func(t *testing.T) {
|
|
// 1. Add Text Document
|
|
addParams := &api.AddTextParams{
|
|
CollectionID: collectionID,
|
|
Text: "This is a comprehensive test of the text document lifecycle including creation, retrieval, and removal.",
|
|
Locale: "en",
|
|
Metadata: map[string]interface{}{
|
|
"title": "Text Lifecycle Test",
|
|
"description": "Full lifecycle integration test",
|
|
},
|
|
Chunking: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.structured",
|
|
},
|
|
Embedding: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.openai",
|
|
OptionID: "text-embedding-3-small",
|
|
},
|
|
AuthScope: map[string]interface{}{
|
|
"__yao_created_by": "integration_test",
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.AddText(ctx, addParams)
|
|
if err != nil {
|
|
t.Skipf("Skipping integration test due to AddText error: %v", err)
|
|
return
|
|
}
|
|
assert.NotNil(t, result)
|
|
t.Logf("1. Created text document: %s", result.DocID)
|
|
|
|
// 2. Get Document
|
|
doc, err := kb.API.GetDocument(ctx, result.DocID, nil)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, doc)
|
|
assert.Equal(t, "Text Lifecycle Test", doc["name"])
|
|
assert.Equal(t, "text", doc["type"])
|
|
assert.Equal(t, "completed", doc["status"])
|
|
t.Logf("2. Retrieved document: name=%v, status=%v", doc["name"], doc["status"])
|
|
|
|
// 3. List Documents
|
|
listFilter := &api.ListDocumentsFilter{
|
|
Page: 1,
|
|
PageSize: 20,
|
|
CollectionID: collectionID,
|
|
Keywords: "Text Lifecycle",
|
|
}
|
|
listResult, err := kb.API.ListDocuments(ctx, listFilter)
|
|
assert.NoError(t, err)
|
|
assert.GreaterOrEqual(t, len(listResult.Data), 1)
|
|
t.Logf("3. Found document in list: %d documents", len(listResult.Data))
|
|
|
|
// 4. Remove Document
|
|
removeParams := &api.RemoveDocumentsParams{
|
|
DocumentIDs: []string{result.DocID},
|
|
}
|
|
removeResult, err := kb.API.RemoveDocuments(ctx, removeParams)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, removeResult)
|
|
t.Logf("4. Removed document: %d deleted", removeResult.DeletedCount)
|
|
|
|
// 5. Verify Removal
|
|
_, err = kb.API.GetDocument(ctx, result.DocID, nil)
|
|
assert.Error(t, err)
|
|
t.Logf("5. Verified document removal")
|
|
|
|
t.Logf("✅ Full text lifecycle test completed successfully")
|
|
})
|
|
}
|