yao/openapi/tests/kb/utils_test.go
Max 91d67d3b54 Enhance collection management and segment handling in the API
- Added GetCollection endpoint to retrieve collections by ID, including error handling for missing IDs and uninitialized knowledge base.
- Updated CreateCollectionConfig to use more descriptive field names for embedding provider and option IDs.
- Refactored UpdateSegments function to improve validation and error handling, ensuring document IDs and segment texts are properly checked.
- Introduced new request structures for updating votes, scores, and weights for segments, enhancing the API's capabilities.
- Removed the outdated vote handling file to streamline the codebase.
2025-08-20 18:25:28 +08:00

517 lines
12 KiB
Go

package openapi_test
import (
"testing"
"github.com/yaoapp/gou/graphrag/types"
kbtypes "github.com/yaoapp/yao/kb/types"
"github.com/yaoapp/yao/openapi/kb"
"github.com/yaoapp/yao/openapi/tests/testutils"
)
func TestProviderConfig_Validation(t *testing.T) {
tests := []struct {
name string
config *kb.ProviderConfig
expectError bool
}{
{
name: "valid provider config with option_id",
config: &kb.ProviderConfig{
ProviderID: "test_provider",
OptionID: "test_option",
},
expectError: false,
},
{
name: "valid provider config with direct option",
config: &kb.ProviderConfig{
ProviderID: "test_provider",
Option: &kbtypes.ProviderOption{
Label: "Test Option",
Value: "test",
Description: "Test description",
Properties: map[string]interface{}{"key": "value"},
},
},
expectError: false,
},
{
name: "invalid provider config - empty provider_id",
config: &kb.ProviderConfig{
ProviderID: "",
OptionID: "test_option",
},
expectError: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if tt.config.ProviderID == "" && !tt.expectError {
t.Errorf("Expected validation to fail for empty provider_id")
}
})
}
}
func TestBaseUpsertRequest_Validate(t *testing.T) {
tests := []struct {
name string
request *kb.BaseUpsertRequest
expectError bool
errorMsg string
}{
{
name: "valid base request",
request: &kb.BaseUpsertRequest{
CollectionID: "test_collection",
Chunking: &kb.ProviderConfig{
ProviderID: "chunking_provider",
OptionID: "default",
},
Embedding: &kb.ProviderConfig{
ProviderID: "embedding_provider",
OptionID: "default",
},
},
expectError: false,
},
{
name: "missing collection_id",
request: &kb.BaseUpsertRequest{
Chunking: &kb.ProviderConfig{
ProviderID: "chunking_provider",
},
Embedding: &kb.ProviderConfig{
ProviderID: "embedding_provider",
},
},
expectError: true,
errorMsg: "collection_id is required",
},
{
name: "missing chunking provider",
request: &kb.BaseUpsertRequest{
CollectionID: "test_collection",
Embedding: &kb.ProviderConfig{
ProviderID: "embedding_provider",
},
},
expectError: true,
errorMsg: "chunking provider is required",
},
{
name: "missing embedding provider",
request: &kb.BaseUpsertRequest{
CollectionID: "test_collection",
Chunking: &kb.ProviderConfig{
ProviderID: "chunking_provider",
},
},
expectError: true,
errorMsg: "embedding provider is required",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
err := tt.request.Validate()
if tt.expectError {
if err == nil {
t.Errorf("Expected error but got none")
} else if err.Error() != tt.errorMsg {
t.Errorf("Expected error message '%s', got '%s'", tt.errorMsg, err.Error())
}
} else {
if err != nil {
t.Errorf("Expected no error but got: %v", err)
}
}
})
}
}
func TestAddFileRequest_Validate(t *testing.T) {
tests := []struct {
name string
request *kb.AddFileRequest
expectError bool
errorMsg string
}{
{
name: "valid add file request",
request: &kb.AddFileRequest{
BaseUpsertRequest: kb.BaseUpsertRequest{
CollectionID: "test_collection",
Chunking: &kb.ProviderConfig{
ProviderID: "chunking_provider",
},
Embedding: &kb.ProviderConfig{
ProviderID: "embedding_provider",
},
},
FileID: "test_file_123",
},
expectError: false,
},
{
name: "missing file_id",
request: &kb.AddFileRequest{
BaseUpsertRequest: kb.BaseUpsertRequest{
CollectionID: "test_collection",
Chunking: &kb.ProviderConfig{
ProviderID: "chunking_provider",
},
Embedding: &kb.ProviderConfig{
ProviderID: "embedding_provider",
},
},
FileID: "",
},
expectError: true,
errorMsg: "file_id is required",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
err := tt.request.Validate()
if tt.expectError {
if err == nil {
t.Errorf("Expected error but got none")
} else if err.Error() != tt.errorMsg {
t.Errorf("Expected error message '%s', got '%s'", tt.errorMsg, err.Error())
}
} else {
if err != nil {
t.Errorf("Expected no error but got: %v", err)
}
}
})
}
}
func TestAddTextRequest_Validate(t *testing.T) {
tests := []struct {
name string
request *kb.AddTextRequest
expectError bool
errorMsg string
}{
{
name: "valid add text request",
request: &kb.AddTextRequest{
BaseUpsertRequest: kb.BaseUpsertRequest{
CollectionID: "test_collection",
Chunking: &kb.ProviderConfig{
ProviderID: "chunking_provider",
},
Embedding: &kb.ProviderConfig{
ProviderID: "embedding_provider",
},
},
Text: "This is test text content",
},
expectError: false,
},
{
name: "missing text",
request: &kb.AddTextRequest{
BaseUpsertRequest: kb.BaseUpsertRequest{
CollectionID: "test_collection",
Chunking: &kb.ProviderConfig{
ProviderID: "chunking_provider",
},
Embedding: &kb.ProviderConfig{
ProviderID: "embedding_provider",
},
},
Text: "",
},
expectError: true,
errorMsg: "text is required",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
err := tt.request.Validate()
if tt.expectError {
if err == nil {
t.Errorf("Expected error but got none")
} else if err.Error() != tt.errorMsg {
t.Errorf("Expected error message '%s', got '%s'", tt.errorMsg, err.Error())
}
} else {
if err != nil {
t.Errorf("Expected no error but got: %v", err)
}
}
})
}
}
func TestAddURLRequest_Validate(t *testing.T) {
tests := []struct {
name string
request *kb.AddURLRequest
expectError bool
errorMsg string
}{
{
name: "valid add URL request",
request: &kb.AddURLRequest{
BaseUpsertRequest: kb.BaseUpsertRequest{
CollectionID: "test_collection",
Chunking: &kb.ProviderConfig{
ProviderID: "chunking_provider",
},
Embedding: &kb.ProviderConfig{
ProviderID: "embedding_provider",
},
},
URL: "https://example.com/document",
},
expectError: false,
},
{
name: "missing URL",
request: &kb.AddURLRequest{
BaseUpsertRequest: kb.BaseUpsertRequest{
CollectionID: "test_collection",
Chunking: &kb.ProviderConfig{
ProviderID: "chunking_provider",
},
Embedding: &kb.ProviderConfig{
ProviderID: "embedding_provider",
},
},
URL: "",
},
expectError: true,
errorMsg: "url is required",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
err := tt.request.Validate()
if tt.expectError {
if err == nil {
t.Errorf("Expected error but got none")
} else if err.Error() != tt.errorMsg {
t.Errorf("Expected error message '%s', got '%s'", tt.errorMsg, err.Error())
}
} else {
if err != nil {
t.Errorf("Expected no error but got: %v", err)
}
}
})
}
}
func TestAddSegmentsRequest_Validate(t *testing.T) {
tests := []struct {
name string
request *kb.AddSegmentsRequest
expectError bool
errorMsg string
}{
{
name: "valid add segments request",
request: &kb.AddSegmentsRequest{
BaseUpsertRequest: kb.BaseUpsertRequest{
CollectionID: "test_collection",
DocID: "test_doc_123",
Chunking: &kb.ProviderConfig{
ProviderID: "chunking_provider",
},
Embedding: &kb.ProviderConfig{
ProviderID: "embedding_provider",
},
},
SegmentTexts: []types.SegmentText{
{Text: "First segment"},
{Text: "Second segment"},
},
},
expectError: false,
},
{
name: "missing segment_texts",
request: &kb.AddSegmentsRequest{
BaseUpsertRequest: kb.BaseUpsertRequest{
CollectionID: "test_collection",
DocID: "test_doc_123",
Chunking: &kb.ProviderConfig{
ProviderID: "chunking_provider",
},
Embedding: &kb.ProviderConfig{
ProviderID: "embedding_provider",
},
},
SegmentTexts: []types.SegmentText{},
},
expectError: true,
errorMsg: "segment_texts is required",
},
{
name: "missing doc_id",
request: &kb.AddSegmentsRequest{
BaseUpsertRequest: kb.BaseUpsertRequest{
CollectionID: "test_collection",
Chunking: &kb.ProviderConfig{
ProviderID: "chunking_provider",
},
Embedding: &kb.ProviderConfig{
ProviderID: "embedding_provider",
},
},
SegmentTexts: []types.SegmentText{
{Text: "Test segment"},
},
},
expectError: true,
errorMsg: "doc_id is required for AddSegments operation",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
err := tt.request.Validate()
if tt.expectError {
if err == nil {
t.Errorf("Expected error but got none")
} else if err.Error() != tt.errorMsg {
t.Errorf("Expected error message '%s', got '%s'", tt.errorMsg, err.Error())
}
} else {
if err != nil {
t.Errorf("Expected no error but got: %v", err)
}
}
})
}
}
func TestUpdateSegmentsRequest_Structure(t *testing.T) {
tests := []struct {
name string
request *kb.UpdateSegmentsRequest
expectValid bool
}{
{
name: "valid update segments request",
request: &kb.UpdateSegmentsRequest{
SegmentTexts: []types.SegmentText{
{ID: "segment_1", Text: "Updated segment"},
},
},
expectValid: true,
},
{
name: "empty segment_texts",
request: &kb.UpdateSegmentsRequest{
SegmentTexts: []types.SegmentText{},
},
expectValid: false,
},
{
name: "multiple segments",
request: &kb.UpdateSegmentsRequest{
SegmentTexts: []types.SegmentText{
{ID: "segment_1", Text: "First updated segment"},
{ID: "segment_2", Text: "Second updated segment"},
},
},
expectValid: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
// Test basic structure validation
if tt.expectValid {
if len(tt.request.SegmentTexts) == 0 {
t.Errorf("Expected valid request to have segment_texts")
}
} else {
if len(tt.request.SegmentTexts) > 0 {
t.Errorf("Expected invalid request to have empty segment_texts")
}
}
})
}
}
func TestToUpsertOptions_WithEnvironment(t *testing.T) {
// Initialize test environment
serverURL := testutils.Prepare(t)
defer testutils.Clean()
t.Logf("Test server running at: %s", serverURL)
// Test with proper environment setup and direct options to avoid provider lookup
request := &kb.BaseUpsertRequest{
CollectionID: "test_collection",
DocID: "test_doc",
Metadata: map[string]interface{}{"source": "test"},
Chunking: &kb.ProviderConfig{
ProviderID: "chunking_provider",
Option: &kbtypes.ProviderOption{
Label: "Test Chunking",
Value: "test",
Description: "Test chunking option",
Properties: map[string]interface{}{"chunk_size": 1000},
},
},
Embedding: &kb.ProviderConfig{
ProviderID: "embedding_provider",
Option: &kbtypes.ProviderOption{
Label: "Test Embedding",
Value: "test",
Description: "Test embedding option",
Properties: map[string]interface{}{"model": "test-model"},
},
},
}
t.Run("no parameters", func(t *testing.T) {
options, err := request.ToUpsertOptions()
if err != nil {
t.Logf("ToUpsertOptions failed (expected with test environment): %v", err)
// This is expected to fail in test environment due to missing actual providers
// But we can verify the basic structure is being built correctly
return
}
// If it doesn't fail, verify the basic structure
if options.CollectionID != "test_collection" {
t.Errorf("Expected CollectionID to be 'test_collection', got '%s'", options.CollectionID)
}
if options.DocID != "test_doc" {
t.Errorf("Expected DocID to be 'test_doc', got '%s'", options.DocID)
}
})
t.Run("with filename and contentType", func(t *testing.T) {
options, err := request.ToUpsertOptions("test.pdf", "application/pdf")
if err != nil {
t.Logf("ToUpsertOptions with file info failed (expected with test environment): %v", err)
// This is expected to fail in test environment due to missing actual providers
return
}
// If it doesn't fail, verify the basic structure
if options.CollectionID != "test_collection" {
t.Errorf("Expected CollectionID to be 'test_collection', got '%s'", options.CollectionID)
}
})
}