yao/openapi/tests/kb/utils_test.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_Validate(t *testing.T) {
	tests := []struct {
		name        string
		request     *kb.UpdateSegmentsRequest
		expectError bool
		errorMsg    string
	}{
		{
			name: "valid update segments request",
			request: &kb.UpdateSegmentsRequest{
				BaseUpsertRequest: kb.BaseUpsertRequest{
					CollectionID: "test_collection",
					Chunking: &kb.ProviderConfig{
						ProviderID: "chunking_provider",
					},
					Embedding: &kb.ProviderConfig{
						ProviderID: "embedding_provider",
					},
				},
				SegmentTexts: []types.SegmentText{
					{Text: "Updated segment"},
				},
			},
			expectError: false,
		},
		{
			name: "missing segment_texts",
			request: &kb.UpdateSegmentsRequest{
				BaseUpsertRequest: kb.BaseUpsertRequest{
					CollectionID: "test_collection",
					Chunking: &kb.ProviderConfig{
						ProviderID: "chunking_provider",
					},
					Embedding: &kb.ProviderConfig{
						ProviderID: "embedding_provider",
					},
				},
				SegmentTexts: []types.SegmentText{},
			},
			expectError: true,
			errorMsg:    "segment_texts 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 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)
		}
	})
}