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.
585 lines
16 KiB
Go
585 lines
16 KiB
Go
package api_test
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"mime/multipart"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
graphragtypes "github.com/yaoapp/gou/graphrag/types"
|
|
"github.com/yaoapp/yao/attachment"
|
|
"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/...
|
|
|
|
// createTestCollectionForFile is a helper to create a test collection for file tests
|
|
func createTestCollectionForFile(t *testing.T, ctx context.Context) string {
|
|
if kb.API == nil {
|
|
t.Skip("KB API not initialized")
|
|
}
|
|
|
|
collectionID := fmt.Sprintf("test_file_%d", time.Now().UnixNano())
|
|
|
|
params := &api.CreateCollectionParams{
|
|
ID: collectionID,
|
|
Metadata: map[string]interface{}{
|
|
"name": "Test File Collection",
|
|
"description": "Collection for AddFile 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
|
|
}
|
|
|
|
// cleanupTestCollectionForFile removes a test collection
|
|
func cleanupTestCollectionForFile(ctx context.Context, collectionID string) {
|
|
if kb.API != nil {
|
|
_, _ = kb.API.RemoveCollection(ctx, collectionID)
|
|
}
|
|
}
|
|
|
|
// ========== AddFile Tests ==========
|
|
// Note: Full AddFile tests require actual files to be uploaded via attachment manager
|
|
// These tests verify parameter validation and error handling
|
|
|
|
func TestAddFile(t *testing.T) {
|
|
if kb.API == nil {
|
|
t.Skip("KB API not initialized")
|
|
}
|
|
|
|
ctx := context.Background()
|
|
collectionID := createTestCollectionForFile(t, ctx)
|
|
defer cleanupTestCollectionForFile(ctx, collectionID)
|
|
|
|
t.Run("AddFileMissingCollectionID", func(t *testing.T) {
|
|
params := &api.AddFileParams{
|
|
FileID: "some_file_id",
|
|
Chunking: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.structured",
|
|
},
|
|
Embedding: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.openai",
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.AddFile(ctx, params)
|
|
assert.Error(t, err)
|
|
assert.Nil(t, result)
|
|
assert.Contains(t, err.Error(), "collection_id is required")
|
|
})
|
|
|
|
t.Run("AddFileMissingFileID", func(t *testing.T) {
|
|
params := &api.AddFileParams{
|
|
CollectionID: collectionID,
|
|
Chunking: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.structured",
|
|
},
|
|
Embedding: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.openai",
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.AddFile(ctx, params)
|
|
assert.Error(t, err)
|
|
assert.Nil(t, result)
|
|
assert.Contains(t, err.Error(), "file_id is required")
|
|
})
|
|
|
|
t.Run("AddFileMissingChunking", func(t *testing.T) {
|
|
params := &api.AddFileParams{
|
|
CollectionID: collectionID,
|
|
FileID: "some_file_id",
|
|
Embedding: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.openai",
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.AddFile(ctx, params)
|
|
assert.Error(t, err)
|
|
assert.Nil(t, result)
|
|
assert.Contains(t, err.Error(), "chunking configuration is required")
|
|
})
|
|
|
|
t.Run("AddFileMissingEmbedding", func(t *testing.T) {
|
|
params := &api.AddFileParams{
|
|
CollectionID: collectionID,
|
|
FileID: "some_file_id",
|
|
Chunking: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.structured",
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.AddFile(ctx, params)
|
|
assert.Error(t, err)
|
|
assert.Nil(t, result)
|
|
assert.Contains(t, err.Error(), "embedding configuration is required")
|
|
})
|
|
|
|
t.Run("AddFileInvalidUploader", func(t *testing.T) {
|
|
params := &api.AddFileParams{
|
|
CollectionID: collectionID,
|
|
FileID: "some_file_id",
|
|
Uploader: "invalid_uploader",
|
|
Chunking: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.structured",
|
|
},
|
|
Embedding: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.openai",
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.AddFile(ctx, params)
|
|
assert.Error(t, err)
|
|
assert.Nil(t, result)
|
|
assert.Contains(t, err.Error(), "invalid uploader")
|
|
})
|
|
|
|
t.Run("AddFileNotFound", func(t *testing.T) {
|
|
params := &api.AddFileParams{
|
|
CollectionID: collectionID,
|
|
FileID: "nonexistent_file_id",
|
|
Chunking: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.structured",
|
|
},
|
|
Embedding: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.openai",
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.AddFile(ctx, params)
|
|
assert.Error(t, err)
|
|
assert.Nil(t, result)
|
|
// Error could be "file not found" or "invalid uploader" depending on environment
|
|
assert.True(t, err != nil, "Expected an error")
|
|
})
|
|
}
|
|
|
|
// ========== AddFileAsync Tests ==========
|
|
|
|
func TestAddFileAsync(t *testing.T) {
|
|
if kb.API == nil {
|
|
t.Skip("KB API not initialized")
|
|
}
|
|
|
|
ctx := context.Background()
|
|
collectionID := createTestCollectionForFile(t, ctx)
|
|
defer cleanupTestCollectionForFile(ctx, collectionID)
|
|
|
|
t.Run("AddFileAsyncMissingCollectionID", func(t *testing.T) {
|
|
params := &api.AddFileParams{
|
|
FileID: "some_file_id",
|
|
Chunking: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.structured",
|
|
},
|
|
Embedding: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.openai",
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.AddFileAsync(ctx, params)
|
|
assert.Error(t, err)
|
|
assert.Nil(t, result)
|
|
assert.Contains(t, err.Error(), "collection_id is required")
|
|
})
|
|
|
|
t.Run("AddFileAsyncMissingFileID", func(t *testing.T) {
|
|
params := &api.AddFileParams{
|
|
CollectionID: collectionID,
|
|
Chunking: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.structured",
|
|
},
|
|
Embedding: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.openai",
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.AddFileAsync(ctx, params)
|
|
assert.Error(t, err)
|
|
assert.Nil(t, result)
|
|
assert.Contains(t, err.Error(), "file_id is required")
|
|
})
|
|
|
|
t.Run("AddFileAsyncMissingChunking", func(t *testing.T) {
|
|
params := &api.AddFileParams{
|
|
CollectionID: collectionID,
|
|
FileID: "some_file_id",
|
|
Embedding: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.openai",
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.AddFileAsync(ctx, params)
|
|
assert.Error(t, err)
|
|
assert.Nil(t, result)
|
|
assert.Contains(t, err.Error(), "chunking configuration is required")
|
|
})
|
|
|
|
t.Run("AddFileAsyncMissingEmbedding", func(t *testing.T) {
|
|
params := &api.AddFileParams{
|
|
CollectionID: collectionID,
|
|
FileID: "some_file_id",
|
|
Chunking: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.structured",
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.AddFileAsync(ctx, params)
|
|
assert.Error(t, err)
|
|
assert.Nil(t, result)
|
|
assert.Contains(t, err.Error(), "embedding configuration is required")
|
|
})
|
|
|
|
t.Run("AddFileAsyncInvalidUploader", func(t *testing.T) {
|
|
params := &api.AddFileParams{
|
|
CollectionID: collectionID,
|
|
FileID: "some_file_id",
|
|
Uploader: "invalid_uploader",
|
|
Chunking: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.structured",
|
|
},
|
|
Embedding: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.openai",
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.AddFileAsync(ctx, params)
|
|
assert.Error(t, err)
|
|
assert.Nil(t, result)
|
|
assert.Contains(t, err.Error(), "invalid uploader")
|
|
})
|
|
|
|
t.Run("AddFileAsyncNotFound", func(t *testing.T) {
|
|
params := &api.AddFileParams{
|
|
CollectionID: collectionID,
|
|
FileID: "nonexistent_file_id",
|
|
Chunking: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.structured",
|
|
},
|
|
Embedding: &api.ProviderConfigParams{
|
|
ProviderID: "__yao.openai",
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.AddFileAsync(ctx, params)
|
|
assert.Error(t, err)
|
|
assert.Nil(t, result)
|
|
// Error could be "file not found" or "invalid uploader" depending on environment
|
|
assert.True(t, err != nil, "Expected an error")
|
|
})
|
|
}
|
|
|
|
// ========== AddFile with Real File Tests ==========
|
|
|
|
// getTestUploader returns the uploader name and manager for testing
|
|
func getTestUploader(t *testing.T) (string, *attachment.Manager) {
|
|
// Try __yao.attachment first (system uploader)
|
|
if manager, ok := attachment.Managers["__yao.attachment"]; ok {
|
|
return "__yao.attachment", manager
|
|
}
|
|
// Try local manager
|
|
if manager, ok := attachment.Managers["local"]; ok {
|
|
return "local", manager
|
|
}
|
|
// List available managers for debugging
|
|
var available []string
|
|
for name := range attachment.Managers {
|
|
available = append(available, name)
|
|
}
|
|
t.Fatalf("No attachment manager available. Available managers: %v", available)
|
|
return "", nil
|
|
}
|
|
|
|
// uploadTestFile uploads a test file using the attachment manager and returns the file ID
|
|
func uploadTestFile(t *testing.T, ctx context.Context, filename, content string) string {
|
|
_, manager := getTestUploader(t)
|
|
|
|
// Create file header
|
|
fileHeader := &attachment.FileHeader{
|
|
FileHeader: &multipart.FileHeader{
|
|
Filename: filename,
|
|
Size: int64(len(content)),
|
|
Header: make(map[string][]string),
|
|
},
|
|
}
|
|
fileHeader.Header.Set("Content-Type", "text/plain")
|
|
|
|
// Upload the file
|
|
reader := strings.NewReader(content)
|
|
file, err := manager.Upload(ctx, fileHeader, reader, attachment.UploadOption{})
|
|
if err != nil {
|
|
t.Fatalf("Failed to upload test file: %v", err)
|
|
}
|
|
|
|
t.Logf("Uploaded test file: %s (ID: %s)", filename, file.ID)
|
|
return file.ID
|
|
}
|
|
|
|
// cleanupTestFile removes a test file
|
|
func cleanupTestFile(ctx context.Context, t *testing.T, fileID string) {
|
|
_, manager := getTestUploader(t)
|
|
_ = manager.Delete(ctx, fileID)
|
|
}
|
|
|
|
func TestAddFileWithRealFile(t *testing.T) {
|
|
if kb.API == nil {
|
|
t.Skip("KB API not initialized")
|
|
}
|
|
|
|
ctx := context.Background()
|
|
collectionID := createTestCollectionForFile(t, ctx)
|
|
defer cleanupTestCollectionForFile(ctx, collectionID)
|
|
|
|
// Get the uploader name
|
|
uploaderName, _ := getTestUploader(t)
|
|
|
|
// Upload a test file
|
|
testContent := `This is a test document for the knowledge base.
|
|
It contains content to test the file processing functionality.`
|
|
|
|
fileID := uploadTestFile(t, ctx, "test_document.txt", testContent)
|
|
defer cleanupTestFile(ctx, t, fileID)
|
|
|
|
t.Run("AddFileSuccess", func(t *testing.T) {
|
|
params := &api.AddFileParams{
|
|
CollectionID: collectionID,
|
|
FileID: fileID,
|
|
Uploader: uploaderName,
|
|
Locale: "en",
|
|
Metadata: map[string]interface{}{
|
|
"description": "A test file 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.AddFile(ctx, params)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, result)
|
|
if result != nil {
|
|
assert.Equal(t, collectionID, result.CollectionID)
|
|
assert.NotEmpty(t, result.DocID)
|
|
assert.Equal(t, fileID, result.FileID)
|
|
assert.Contains(t, result.Message, "successfully")
|
|
t.Logf("Added file document: %s", result.DocID)
|
|
}
|
|
|
|
// Verify document was created
|
|
if result != nil {
|
|
doc, err := kb.API.GetDocument(ctx, result.DocID, nil)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, doc)
|
|
assert.Equal(t, "file", doc["type"])
|
|
assert.Equal(t, "completed", doc["status"])
|
|
t.Logf("✅ File Document verified: type=%v, status=%v", doc["type"], doc["status"])
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestAddFileAsyncWithRealFile(t *testing.T) {
|
|
if kb.API == nil {
|
|
t.Skip("KB API not initialized")
|
|
}
|
|
|
|
ctx := context.Background()
|
|
collectionID := createTestCollectionForFile(t, ctx)
|
|
defer cleanupTestCollectionForFile(ctx, collectionID)
|
|
|
|
// Get the uploader name
|
|
uploaderName, _ := getTestUploader(t)
|
|
|
|
// Upload a test file for async processing
|
|
testContent := `Async test document content.
|
|
|
|
This document will be processed asynchronously.
|
|
|
|
The job system should handle the processing in the background.`
|
|
|
|
fileID := uploadTestFile(t, ctx, "async_test_document.txt", testContent)
|
|
defer cleanupTestFile(ctx, t, fileID)
|
|
|
|
t.Run("AddFileAsyncSuccess", func(t *testing.T) {
|
|
params := &api.AddFileParams{
|
|
CollectionID: collectionID,
|
|
FileID: fileID,
|
|
Uploader: uploaderName,
|
|
Locale: "en",
|
|
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 File Job",
|
|
Description: "Testing async file processing",
|
|
Category: "Test",
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.AddFileAsync(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 file 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, "file", doc["type"])
|
|
assert.Equal(t, result.JobID, doc["job_id"])
|
|
t.Logf("✅ Async file 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))
|
|
assert.Equal(t, "completed", finalStatus, "Job should complete successfully")
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestAddFileIntegration(t *testing.T) {
|
|
if kb.API == nil {
|
|
t.Skip("KB API not initialized")
|
|
}
|
|
|
|
ctx := context.Background()
|
|
collectionID := createTestCollectionForFile(t, ctx)
|
|
defer cleanupTestCollectionForFile(ctx, collectionID)
|
|
|
|
// Get the uploader name
|
|
uploaderName, _ := getTestUploader(t)
|
|
|
|
t.Run("FullFileLifecycle", func(t *testing.T) {
|
|
// Upload a test file
|
|
testContent := `Integration test document.
|
|
|
|
This document tests the full lifecycle of file processing:
|
|
1. Upload file
|
|
2. Add to knowledge base
|
|
3. Verify document creation
|
|
4. List documents
|
|
5. Remove document
|
|
|
|
End of test content.`
|
|
|
|
fileID := uploadTestFile(t, ctx, "lifecycle_test.txt", testContent)
|
|
defer cleanupTestFile(ctx, t, fileID)
|
|
|
|
// 1. Add File Document
|
|
addParams := &api.AddFileParams{
|
|
CollectionID: collectionID,
|
|
FileID: fileID,
|
|
Uploader: uploaderName,
|
|
Locale: "en",
|
|
Metadata: map[string]interface{}{
|
|
"title": "File Lifecycle Test",
|
|
"description": "Full lifecycle integration test",
|
|
},
|
|
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": "integration_test",
|
|
},
|
|
}
|
|
|
|
result, err := kb.API.AddFile(ctx, addParams)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, result)
|
|
if result == nil {
|
|
t.Fatalf("Failed to create file document: result is nil")
|
|
}
|
|
t.Logf("1. Created file 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, "file", doc["type"])
|
|
assert.Equal(t, "completed", doc["status"])
|
|
t.Logf("2. Retrieved document: name=%v, type=%v, status=%v", doc["name"], doc["type"], doc["status"])
|
|
|
|
// 3. List Documents
|
|
listFilter := &api.ListDocumentsFilter{
|
|
Page: 1,
|
|
PageSize: 20,
|
|
CollectionID: collectionID,
|
|
}
|
|
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 file lifecycle test completed successfully")
|
|
})
|
|
}
|