From bb64eef20d58f0aee36545541ef78d57fb8f2f5d Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 24 Jul 2025 17:08:58 +0800 Subject: [PATCH] Implement environment variable resolution in configuration parsing and enhance collection management API - Added a new test for resolving environment variables during configuration parsing, ensuring that environment variables are correctly substituted in the configuration. - Refactored the configuration handling to resolve environment variables immediately after parsing, improving the clarity and usability of the configuration structure. - Enhanced the collection management API by standardizing error responses and success responses using a custom response structure, improving consistency across endpoints. - Implemented cleanup logic for test collections to ensure proper resource management during testing. --- kb/config.go | 125 +++++- kb/config_test.go | 76 ++++ openapi/kb/collection.go | 159 +++++++- openapi/tests/kb/collection_test.go | 566 +++++++++++++++++++++++++++ openapi/tests/testutils/testutils.go | 85 +++- 5 files changed, 981 insertions(+), 30 deletions(-) create mode 100644 openapi/tests/kb/collection_test.go diff --git a/kb/config.go b/kb/config.go index 731b8b1c..f5fcf701 100644 --- a/kb/config.go +++ b/kb/config.go @@ -116,14 +116,31 @@ func (c *Config) createGraphStore() (types.GraphStore, error) { // toVectorStoreConfig converts the vector config to VectorStoreConfig func (c *Config) toVectorStoreConfig() (types.VectorStoreConfig, error) { - // Parse environment variables in config - resolvedConfig, err := c.resolveEnvVars(c.Vector.Config) - if err != nil { - return types.VectorStoreConfig{}, err + // Environment variables are already resolved during parsing + configCopy := make(map[string]interface{}) + for k, v := range c.Vector.Config { + configCopy[k] = v } - // Convert resolved config to types.VectorStoreConfig via JSON - jsonData, err := json.Marshal(resolvedConfig) + // Ensure host and port are in ExtraParams for Qdrant + if _, exists := configCopy["extra_params"]; !exists { + configCopy["extra_params"] = make(map[string]interface{}) + } + + extraParams := configCopy["extra_params"].(map[string]interface{}) + + // Map host field + if host, exists := configCopy["host"]; exists { + extraParams["host"] = host + } + + // Map port field + if port, exists := configCopy["port"]; exists { + extraParams["port"] = port + } + + // Convert config to types.VectorStoreConfig via JSON + jsonData, err := json.Marshal(configCopy) if err != nil { return types.VectorStoreConfig{}, err } @@ -138,14 +155,39 @@ func (c *Config) toVectorStoreConfig() (types.VectorStoreConfig, error) { // toGraphStoreConfig converts the graph config to GraphStoreConfig func (c *Config) toGraphStoreConfig() (types.GraphStoreConfig, error) { - // Parse environment variables in config - resolvedConfig, err := c.resolveEnvVars(c.Graph.Config) - if err != nil { - return types.GraphStoreConfig{}, err + // Environment variables are already resolved during parsing + configCopy := make(map[string]interface{}) + for k, v := range c.Graph.Config { + configCopy[k] = v } - // Convert resolved config to types.GraphStoreConfig via JSON - jsonData, err := json.Marshal(resolvedConfig) + // Map field names to match GraphStoreConfig structure + if url, exists := configCopy["url"]; exists { + configCopy["database_url"] = url + delete(configCopy, "url") // Remove the original field + } + + // Ensure DriverConfig exists and map username/password into it + if _, exists := configCopy["driver_config"]; !exists { + configCopy["driver_config"] = make(map[string]interface{}) + } + + driverConfig := configCopy["driver_config"].(map[string]interface{}) + + // Map username field to DriverConfig + if username, exists := configCopy["username"]; exists { + driverConfig["username"] = username + delete(configCopy, "username") // Remove from top level + } + + // Map password field to DriverConfig + if password, exists := configCopy["password"]; exists { + driverConfig["password"] = password + delete(configCopy, "password") // Remove from top level + } + + // Convert config to types.GraphStoreConfig via JSON + jsonData, err := json.Marshal(configCopy) if err != nil { return types.GraphStoreConfig{}, err } @@ -200,6 +242,58 @@ func (c *Config) parseEnvVar(value string) string { }) } +// resolveAllEnvVars resolves environment variables in all configuration sections +func (c *Config) resolveAllEnvVars() error { + // Resolve Vector config + if c.Vector.Config != nil { + resolved, err := c.resolveEnvVars(c.Vector.Config) + if err != nil { + return err + } + c.Vector.Config = resolved + } + + // Resolve Graph config + if c.Graph != nil && c.Graph.Config != nil { + resolved, err := c.resolveEnvVars(c.Graph.Config) + if err != nil { + return err + } + c.Graph.Config = resolved + } + + // Resolve Provider options (if they contain env vars) + if err := c.resolveProviderEnvVars(); err != nil { + return err + } + + return nil +} + +// resolveProviderEnvVars resolves environment variables in provider configurations +func (c *Config) resolveProviderEnvVars() error { + providerLists := [][]*Provider{ + c.Chunkings, c.Embeddings, c.Converters, c.Extractors, + c.Fetchers, c.Searchers, c.Rerankers, c.Votes, c.Weights, c.Scores, + } + + for _, providers := range providerLists { + for _, provider := range providers { + for _, option := range provider.Options { + if option.Properties != nil { + resolved, err := c.resolveEnvVars(option.Properties) + if err != nil { + return err + } + option.Properties = resolved + } + } + } + } + + return nil +} + // UnmarshalJSON implements json.Unmarshaler interface func (c *Config) UnmarshalJSON(data []byte) error { // Use alias type to avoid infinite recursion @@ -208,7 +302,12 @@ func (c *Config) UnmarshalJSON(data []byte) error { return err } - // Compute features after parsing + // Resolve environment variables immediately after parsing + if err := c.resolveAllEnvVars(); err != nil { + return err + } + + // Compute features after parsing and resolving env vars c.Features = c.ComputeFeatures() return nil diff --git a/kb/config_test.go b/kb/config_test.go index bb818fa4..1a28a5d0 100644 --- a/kb/config_test.go +++ b/kb/config_test.go @@ -506,3 +506,79 @@ func TestConfig_ParseEnvVar(t *testing.T) { }) } } + +func TestConfig_ResolveEnvVarsOnParsing(t *testing.T) { + // Set test environment variables + os.Setenv("TEST_VECTOR_HOST", "test-vector-host") + os.Setenv("TEST_GRAPH_URL", "neo4j://test-graph:7687") + os.Setenv("TEST_GRAPH_USER", "test-user") + os.Setenv("TEST_GRAPH_PASS", "test-pass") + defer func() { + os.Unsetenv("TEST_VECTOR_HOST") + os.Unsetenv("TEST_GRAPH_URL") + os.Unsetenv("TEST_GRAPH_USER") + os.Unsetenv("TEST_GRAPH_PASS") + }() + + configJSON := `{ + "vector": { + "driver": "qdrant", + "config": { + "host": "$ENV.TEST_VECTOR_HOST", + "port": 6333 + } + }, + "graph": { + "driver": "neo4j", + "config": { + "url": "$ENV.TEST_GRAPH_URL", + "username": "$ENV.TEST_GRAPH_USER", + "password": "$ENV.TEST_GRAPH_PASS" + } + }, + "chunkings": [ + { + "id": "__yao.structured", + "label": "Document Structure", + "description": "Split text", + "options": [] + } + ], + "embeddings": [ + { + "id": "__yao.openai", + "label": "OpenAI", + "description": "OpenAI embeddings", + "options": [] + } + ] + }` + + // Parse config from JSON + config, err := ParseConfigFromJSON([]byte(configJSON)) + if err != nil { + t.Fatalf("Failed to parse config: %v", err) + } + + // Verify that environment variables are resolved immediately after parsing + if config.Vector.Config["host"] != "test-vector-host" { + t.Errorf("Expected vector host to be resolved to 'test-vector-host', got '%v'", config.Vector.Config["host"]) + } + + if config.Graph.Config["url"] != "neo4j://test-graph:7687" { + t.Errorf("Expected graph URL to be resolved to 'neo4j://test-graph:7687', got '%v'", config.Graph.Config["url"]) + } + + if config.Graph.Config["username"] != "test-user" { + t.Errorf("Expected graph username to be resolved to 'test-user', got '%v'", config.Graph.Config["username"]) + } + + if config.Graph.Config["password"] != "test-pass" { + t.Errorf("Expected graph password to be resolved to 'test-pass', got '%v'", config.Graph.Config["password"]) + } + + // Verify that numeric values remain unchanged (JSON numbers are parsed as float64) + if port, ok := config.Vector.Config["port"].(float64); !ok || port != 6333.0 { + t.Errorf("Expected vector port to remain 6333.0, got %v (type %T)", config.Vector.Config["port"], config.Vector.Config["port"]) + } +} diff --git a/openapi/kb/collection.go b/openapi/kb/collection.go index cf0b223d..e241da8f 100644 --- a/openapi/kb/collection.go +++ b/openapi/kb/collection.go @@ -2,11 +2,11 @@ package kb import ( "fmt" - "net/http" "github.com/gin-gonic/gin" "github.com/yaoapp/gou/graphrag/types" "github.com/yaoapp/yao/kb" + "github.com/yaoapp/yao/openapi/response" ) // Collection Management Handlers @@ -17,19 +17,34 @@ func CreateCollection(c *gin.Context) { // Parse and bind JSON request if err := c.ShouldBindJSON(&req); err != nil { - c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid request format: " + err.Error()}) + // Create a custom error with the same structure but specific message + errorResp := &response.ErrorResponse{ + Code: response.ErrInvalidRequest.Code, + ErrorDescription: "Invalid request format: " + err.Error(), + } + response.RespondWithError(c, response.StatusBadRequest, errorResp) return } // Validate request parameters if err := validateCreateCollectionRequest(&req); err != nil { - c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + // Create a custom error with the same structure but specific message + errorResp := &response.ErrorResponse{ + Code: response.ErrInvalidRequest.Code, + ErrorDescription: err.Error(), + } + response.RespondWithError(c, response.StatusBadRequest, errorResp) return } // Check if kb.Instance is available if kb.Instance == nil { - c.JSON(http.StatusInternalServerError, gin.H{"error": "Knowledge base not initialized"}) + // Create a custom error with the same structure but specific message + errorResp := &response.ErrorResponse{ + Code: response.ErrServerError.Code, + ErrorDescription: "Knowledge base not initialized", + } + response.RespondWithError(c, response.StatusInternalServerError, errorResp) return } @@ -43,36 +58,150 @@ func CreateCollection(c *gin.Context) { // Call the actual CreateCollection method collectionID, err := kb.Instance.CreateCollection(c.Request.Context(), collectionConfig) if err != nil { - c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to create collection: " + err.Error()}) + // Create a custom error with the same structure but specific message + errorResp := &response.ErrorResponse{ + Code: response.ErrServerError.Code, + ErrorDescription: "Failed to create collection: " + err.Error(), + } + response.RespondWithError(c, response.StatusInternalServerError, errorResp) return } - c.JSON(http.StatusCreated, gin.H{ + successData := gin.H{ "message": "Collection created successfully", "collection_id": collectionID, - }) + } + response.RespondWithSuccess(c, response.StatusCreated, successData) } // RemoveCollection removes an existing collection func RemoveCollection(c *gin.Context) { - // TODO: Implement remove collection logic - c.JSON(http.StatusOK, gin.H{"message": "Collection removed"}) + // Get collection ID from URL parameter + collectionID := c.Param("collectionID") + if collectionID == "" { + errorResp := &response.ErrorResponse{ + Code: response.ErrInvalidRequest.Code, + ErrorDescription: "Collection ID is required", + } + response.RespondWithError(c, response.StatusBadRequest, errorResp) + return + } + + // Check if kb.Instance is available + if kb.Instance == nil { + errorResp := &response.ErrorResponse{ + Code: response.ErrServerError.Code, + ErrorDescription: "Knowledge base not initialized", + } + response.RespondWithError(c, response.StatusInternalServerError, errorResp) + return + } + + // Call the actual RemoveCollection method + removed, err := kb.Instance.RemoveCollection(c.Request.Context(), collectionID) + if err != nil { + errorResp := &response.ErrorResponse{ + Code: response.ErrServerError.Code, + ErrorDescription: "Failed to remove collection: " + err.Error(), + } + response.RespondWithError(c, response.StatusInternalServerError, errorResp) + return + } + + if !removed { + errorResp := &response.ErrorResponse{ + Code: response.ErrInvalidRequest.Code, + ErrorDescription: "Collection not found or could not be removed", + } + response.RespondWithError(c, response.StatusNotFound, errorResp) + return + } + + successData := gin.H{ + "message": "Collection removed successfully", + "collection_id": collectionID, + "removed": removed, + } + response.RespondWithSuccess(c, response.StatusOK, successData) } // CollectionExists checks if a collection exists func CollectionExists(c *gin.Context) { - // TODO: Implement collection exists check logic - c.JSON(http.StatusOK, gin.H{"exists": false}) + // Get collection ID from URL parameter + collectionID := c.Param("collectionID") + if collectionID == "" { + errorResp := &response.ErrorResponse{ + Code: response.ErrInvalidRequest.Code, + ErrorDescription: "Collection ID is required", + } + response.RespondWithError(c, response.StatusBadRequest, errorResp) + return + } + + // Check if kb.Instance is available + if kb.Instance == nil { + errorResp := &response.ErrorResponse{ + Code: response.ErrServerError.Code, + ErrorDescription: "Knowledge base not initialized", + } + response.RespondWithError(c, response.StatusInternalServerError, errorResp) + return + } + + // Call the actual CollectionExists method + exists, err := kb.Instance.CollectionExists(c.Request.Context(), collectionID) + if err != nil { + errorResp := &response.ErrorResponse{ + Code: response.ErrServerError.Code, + ErrorDescription: "Failed to check collection existence: " + err.Error(), + } + response.RespondWithError(c, response.StatusInternalServerError, errorResp) + return + } + + successData := gin.H{ + "collection_id": collectionID, + "exists": exists, + } + response.RespondWithSuccess(c, response.StatusOK, successData) } // GetCollections retrieves collections with optional filtering func GetCollections(c *gin.Context) { - collections, err := kb.Instance.GetCollections(c.Request.Context(), map[string]interface{}{}) - if err != nil { - c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + // Check if kb.Instance is available + if kb.Instance == nil { + errorResp := &response.ErrorResponse{ + Code: response.ErrServerError.Code, + ErrorDescription: "Knowledge base not initialized", + } + response.RespondWithError(c, response.StatusInternalServerError, errorResp) return } - c.JSON(http.StatusOK, collections) + + // Build filter from query parameters + filter := make(map[string]interface{}) + + // Extract all query parameters as potential filter conditions + // This allows filtering by any metadata field, e.g.: + // GET /collections?category=documents&status=active + for key, values := range c.Request.URL.Query() { + if len(values) > 0 { + // Use the first value if multiple values are provided + filter[key] = values[0] + } + } + + collections, err := kb.Instance.GetCollections(c.Request.Context(), filter) + if err != nil { + // Create a custom error with the same structure but specific message + errorResp := &response.ErrorResponse{ + Code: response.ErrServerError.Code, + ErrorDescription: err.Error(), + } + response.RespondWithError(c, response.StatusInternalServerError, errorResp) + return + } + response.RespondWithSuccess(c, response.StatusOK, collections) } // CreateCollectionRequest represents the request structure for creating a collection diff --git a/openapi/tests/kb/collection_test.go b/openapi/tests/kb/collection_test.go new file mode 100644 index 00000000..a8980239 --- /dev/null +++ b/openapi/tests/kb/collection_test.go @@ -0,0 +1,566 @@ +package openapi_test + +import ( + "bytes" + "encoding/json" + "fmt" + "net/http" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/yaoapp/yao/openapi" + "github.com/yaoapp/yao/openapi/tests/testutils" +) + +// TestCreateCollection tests the collection creation endpoint +func TestCreateCollection(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + // Get base URL from server config + baseURL := "" + if openapi.Server != nil && openapi.Server.Config != nil { + baseURL = openapi.Server.Config.BaseURL + } + + // Register test client and get token + client := testutils.RegisterTestClient(t, "KB Collection Create Test Client", []string{"https://localhost/callback"}) + defer testutils.CleanupTestClient(t, client.ClientID) + tokenInfo := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile") + + // Generate unique test collection ID + testCollectionID := fmt.Sprintf("test_collection_%d", time.Now().UnixNano()) + + t.Run("CreateCollectionSuccess", func(t *testing.T) { + // Register collection for cleanup + testutils.RegisterTestCollection(testCollectionID) + + // Prepare request body + createData := map[string]interface{}{ + "id": testCollectionID, + "metadata": map[string]interface{}{ + "category": "test", + "created_by": "test_user", + }, + "config": map[string]interface{}{ + "collection_name": testCollectionID + "_vector", // Required: collection name + "embedding_model": "text-embedding-ada-002", + "chunk_size": 1000, + "chunk_overlap": 200, + "index_type": "hnsw", // Required: valid index type + "distance": "cosine", // Required: distance metric + "dimension": 1536, // Required: embedding dimensions (note: singular not plural) + }, + } + + body, err := json.Marshal(createData) + assert.NoError(t, err) + + // Create HTTP request + req, err := http.NewRequest("POST", serverURL+baseURL+"/kb/collections", bytes.NewBuffer(body)) + assert.NoError(t, err) + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + // Make request + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Check response + assert.Equal(t, http.StatusCreated, resp.StatusCode) + + var response map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&response) + assert.NoError(t, err) + + // Debug output for failed requests + if resp.StatusCode != http.StatusCreated { + t.Logf("Expected status 201, got %d", resp.StatusCode) + t.Logf("Response body: %+v", response) + } + + assert.Equal(t, "Collection created successfully", response["message"]) + assert.Equal(t, testCollectionID, response["collection_id"]) + + t.Logf("Successfully created collection: %s", testCollectionID) + }) + + t.Run("CreateCollectionInvalidRequest", func(t *testing.T) { + // Test with missing required fields + invalidData := map[string]interface{}{ + "metadata": map[string]interface{}{ + "category": "test", + }, + // Missing "id" and "config" fields + } + + body, err := json.Marshal(invalidData) + assert.NoError(t, err) + + req, err := http.NewRequest("POST", serverURL+baseURL+"/kb/collections", bytes.NewBuffer(body)) + assert.NoError(t, err) + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) + + var errorResponse map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&errorResponse) + assert.NoError(t, err) + assert.Contains(t, errorResponse, "error") + + t.Logf("Correctly rejected invalid collection creation request") + }) + + t.Run("CreateCollectionMalformedJSON", func(t *testing.T) { + // Test with malformed JSON + malformedJSON := `{"id": "test", "config": malformed}` + + req, err := http.NewRequest("POST", serverURL+baseURL+"/kb/collections", bytes.NewBufferString(malformedJSON)) + assert.NoError(t, err) + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) + + t.Logf("Correctly rejected malformed JSON request") + }) +} + +// TestRemoveCollection tests the collection removal endpoint +func TestRemoveCollection(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + baseURL := "" + if openapi.Server != nil && openapi.Server.Config != nil { + baseURL = openapi.Server.Config.BaseURL + } + + client := testutils.RegisterTestClient(t, "KB Collection Remove Test Client", []string{"https://localhost/callback"}) + defer testutils.CleanupTestClient(t, client.ClientID) + tokenInfo := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile") + + testCollectionID := fmt.Sprintf("test_collection_remove_%d", time.Now().UnixNano()) + + t.Run("RemoveCollectionSuccess", func(t *testing.T) { + // Register collection for cleanup (in case removal fails) + testutils.RegisterTestCollection(testCollectionID) + + // First create a collection to remove + createData := map[string]interface{}{ + "id": testCollectionID, + "metadata": map[string]interface{}{ + "category": "test_remove", + }, + "config": map[string]interface{}{ + "collection_name": testCollectionID + "_vector", + "embedding_model": "text-embedding-ada-002", + "index_type": "hnsw", + "distance": "cosine", + "dimension": 1536, + }, + } + + body, _ := json.Marshal(createData) + createReq, _ := http.NewRequest("POST", serverURL+baseURL+"/kb/collections", bytes.NewBuffer(body)) + createReq.Header.Set("Content-Type", "application/json") + createReq.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + createResp, err := http.DefaultClient.Do(createReq) + assert.NoError(t, err) + defer createResp.Body.Close() + + // Now test removal + req, err := http.NewRequest("DELETE", serverURL+baseURL+"/kb/collections/"+testCollectionID, nil) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Knowledge base should be initialized for this test to be meaningful + if resp.StatusCode == http.StatusInternalServerError { + var errorResponse map[string]interface{} + json.NewDecoder(resp.Body).Decode(&errorResponse) + if errorDescription, ok := errorResponse["error_description"].(string); ok && + strings.Contains(errorDescription, "Knowledge base not initialized") { + t.Skip("Knowledge base not initialized - skipping test (this indicates environment setup issue)") + } + } + + // Expect successful response for collection removal + assert.Equal(t, http.StatusOK, resp.StatusCode, "Should successfully remove collection when KB is initialized") + + var response map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&response) + assert.NoError(t, err) + assert.Contains(t, response, "message") + t.Logf("Successfully removed collection: %s", testCollectionID) + }) + + t.Run("RemoveCollectionMissingID", func(t *testing.T) { + // Test with missing collection ID in path + req, err := http.NewRequest("DELETE", serverURL+baseURL+"/kb/collections/", nil) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // This should result in a 404 or similar error due to route mismatch + assert.NotEqual(t, http.StatusOK, resp.StatusCode) + + t.Logf("Correctly handled request with missing collection ID") + }) +} + +// TestCollectionExists tests the collection existence check endpoint +func TestCollectionExists(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + baseURL := "" + if openapi.Server != nil && openapi.Server.Config != nil { + baseURL = openapi.Server.Config.BaseURL + } + + client := testutils.RegisterTestClient(t, "KB Collection Exists Test Client", []string{"https://localhost/callback"}) + defer testutils.CleanupTestClient(t, client.ClientID) + tokenInfo := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile") + + testCollectionID := fmt.Sprintf("test_collection_exists_%d", time.Now().UnixNano()) + + t.Run("CollectionExistsCheck", func(t *testing.T) { + // Test with a collection ID + req, err := http.NewRequest("GET", serverURL+baseURL+"/kb/collections/"+testCollectionID+"/exists", nil) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Knowledge base should be initialized for this test to be meaningful + if resp.StatusCode == http.StatusInternalServerError { + var errorResponse map[string]interface{} + json.NewDecoder(resp.Body).Decode(&errorResponse) + if errorDescription, ok := errorResponse["error_description"].(string); ok && + strings.Contains(errorDescription, "Knowledge base not initialized") { + t.Skip("Knowledge base not initialized - skipping test (this indicates environment setup issue)") + } + } + + // Expect successful response for collection existence check + assert.Equal(t, http.StatusOK, resp.StatusCode, "Should successfully check collection existence when KB is initialized") + + var response map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&response) + assert.NoError(t, err) + assert.Contains(t, response, "exists") + assert.Contains(t, response, "collection_id") + assert.Equal(t, testCollectionID, response["collection_id"]) + + exists, ok := response["exists"].(bool) + assert.True(t, ok, "exists should be a boolean") + + t.Logf("Collection %s exists: %v", testCollectionID, exists) + }) + + t.Run("CollectionExistsMissingID", func(t *testing.T) { + // Test with missing collection ID in path + req, err := http.NewRequest("GET", serverURL+baseURL+"/kb/collections//exists", nil) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // This should result in a bad request due to empty collection ID + if resp.StatusCode == http.StatusBadRequest { + var errorResponse map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&errorResponse) + assert.NoError(t, err) + assert.Contains(t, errorResponse, "error") + t.Logf("Correctly rejected request with missing collection ID") + } + }) +} + +// TestGetCollections tests the collections listing endpoint +func TestGetCollections(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + baseURL := "" + if openapi.Server != nil && openapi.Server.Config != nil { + baseURL = openapi.Server.Config.BaseURL + } + + client := testutils.RegisterTestClient(t, "KB Collections List Test Client", []string{"https://localhost/callback"}) + defer testutils.CleanupTestClient(t, client.ClientID) + tokenInfo := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile") + + t.Run("GetCollectionsSuccess", func(t *testing.T) { + // Test listing all collections + req, err := http.NewRequest("GET", serverURL+baseURL+"/kb/collections", nil) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Knowledge base should be initialized for this test to be meaningful + if resp.StatusCode == http.StatusInternalServerError { + var errorResponse map[string]interface{} + json.NewDecoder(resp.Body).Decode(&errorResponse) + if errorDescription, ok := errorResponse["error_description"].(string); ok && + strings.Contains(errorDescription, "Knowledge base not initialized") { + t.Skip("Knowledge base not initialized - skipping test (this indicates environment setup issue)") + } + } + + // Expect successful response + assert.Equal(t, http.StatusOK, resp.StatusCode, "Should successfully retrieve collections when KB is initialized") + + var collections []interface{} + err = json.NewDecoder(resp.Body).Decode(&collections) + assert.NoError(t, err) + + t.Logf("Successfully retrieved %d collections", len(collections)) + }) + + t.Run("GetCollectionsWithFilter", func(t *testing.T) { + // Test listing collections with filter parameters + req, err := http.NewRequest("GET", serverURL+baseURL+"/kb/collections?category=documents&status=active", nil) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Knowledge base should be initialized for this test to be meaningful + if resp.StatusCode == http.StatusInternalServerError { + var errorResponse map[string]interface{} + json.NewDecoder(resp.Body).Decode(&errorResponse) + if errorDescription, ok := errorResponse["error_description"].(string); ok && + strings.Contains(errorDescription, "Knowledge base not initialized") { + t.Skip("Knowledge base not initialized - skipping test (this indicates environment setup issue)") + } + } + + // Expect successful response + assert.Equal(t, http.StatusOK, resp.StatusCode, "Should successfully retrieve collections when KB is initialized") + + var collections []interface{} + err = json.NewDecoder(resp.Body).Decode(&collections) + assert.NoError(t, err) + + t.Logf("Successfully retrieved %d filtered collections", len(collections)) + }) + + t.Run("GetCollectionsWithMultipleFilters", func(t *testing.T) { + // Test with multiple filter parameters + req, err := http.NewRequest("GET", serverURL+baseURL+"/kb/collections?category=test&owner=testuser&type=public", nil) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + // Knowledge base should be initialized for this test to be meaningful + if resp.StatusCode == http.StatusInternalServerError { + var errorResponse map[string]interface{} + json.NewDecoder(resp.Body).Decode(&errorResponse) + if errorDescription, ok := errorResponse["error_description"].(string); ok && + strings.Contains(errorDescription, "Knowledge base not initialized") { + t.Skip("Knowledge base not initialized - skipping test (this indicates environment setup issue)") + } + } + + // Expect successful response + assert.Equal(t, http.StatusOK, resp.StatusCode, "Should successfully retrieve collections when KB is initialized") + + var collections []interface{} + err = json.NewDecoder(resp.Body).Decode(&collections) + assert.NoError(t, err) + + t.Logf("Successfully retrieved %d collections with multiple filters", len(collections)) + }) +} + +// TestCollectionEndpointsUnauthorized tests that endpoints return 401 when not authenticated +func TestCollectionEndpointsUnauthorized(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + baseURL := "" + if openapi.Server != nil && openapi.Server.Config != nil { + baseURL = openapi.Server.Config.BaseURL + } + + endpoints := []struct { + method string + path string + body string + }{ + {"GET", "/kb/collections", ""}, + {"POST", "/kb/collections", `{"id":"test","config":{}}`}, + {"DELETE", "/kb/collections/test", ""}, + {"GET", "/kb/collections/test/exists", ""}, + } + + for _, endpoint := range endpoints { + t.Run(fmt.Sprintf("Unauthorized_%s_%s", endpoint.method, endpoint.path), func(t *testing.T) { + var req *http.Request + var err error + + if endpoint.body != "" { + req, err = http.NewRequest(endpoint.method, serverURL+baseURL+endpoint.path, bytes.NewBufferString(endpoint.body)) + req.Header.Set("Content-Type", "application/json") + } else { + req, err = http.NewRequest(endpoint.method, serverURL+baseURL+endpoint.path, nil) + } + assert.NoError(t, err) + + // No Authorization header + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + assert.NotNil(t, resp) + defer resp.Body.Close() + + assert.Equal(t, http.StatusUnauthorized, resp.StatusCode) + + t.Logf("Correctly rejected unauthorized request to %s %s", endpoint.method, endpoint.path) + }) + } +} + +// TestCollectionIntegration tests the full collection lifecycle +func TestCollectionIntegration(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + baseURL := "" + if openapi.Server != nil && openapi.Server.Config != nil { + baseURL = openapi.Server.Config.BaseURL + } + + client := testutils.RegisterTestClient(t, "KB Collection Integration Test Client", []string{"https://localhost/callback"}) + defer testutils.CleanupTestClient(t, client.ClientID) + tokenInfo := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile") + + testCollectionID := fmt.Sprintf("test_integration_%d", time.Now().UnixNano()) + + t.Run("FullCollectionLifecycle", func(t *testing.T) { + // Register collection for cleanup (in case lifecycle test fails) + testutils.RegisterTestCollection(testCollectionID) + + // Step 1: Check that collection doesn't exist initially + req, err := http.NewRequest("GET", serverURL+baseURL+"/kb/collections/"+testCollectionID+"/exists", nil) + assert.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + resp, err := http.DefaultClient.Do(req) + assert.NoError(t, err) + defer resp.Body.Close() + + // Step 2: Create the collection + createData := map[string]interface{}{ + "id": testCollectionID, + "metadata": map[string]interface{}{ + "category": "integration_test", + "purpose": "full_lifecycle_test", + }, + "config": map[string]interface{}{ + "collection_name": testCollectionID + "_vector", + "embedding_model": "text-embedding-ada-002", + "chunk_size": 1000, + "chunk_overlap": 200, + "index_type": "hnsw", + "distance": "cosine", + "dimension": 1536, + }, + } + + body, err := json.Marshal(createData) + assert.NoError(t, err) + + createReq, err := http.NewRequest("POST", serverURL+baseURL+"/kb/collections", bytes.NewBuffer(body)) + assert.NoError(t, err) + createReq.Header.Set("Content-Type", "application/json") + createReq.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + createResp, err := http.DefaultClient.Do(createReq) + assert.NoError(t, err) + defer createResp.Body.Close() + + // Step 3: Verify collection appears in listings (with filter) + listReq, err := http.NewRequest("GET", serverURL+baseURL+"/kb/collections?category=integration_test", nil) + assert.NoError(t, err) + listReq.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + listResp, err := http.DefaultClient.Do(listReq) + assert.NoError(t, err) + defer listResp.Body.Close() + + // Step 4: Check that collection now exists + existsReq, err := http.NewRequest("GET", serverURL+baseURL+"/kb/collections/"+testCollectionID+"/exists", nil) + assert.NoError(t, err) + existsReq.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + existsResp, err := http.DefaultClient.Do(existsReq) + assert.NoError(t, err) + defer existsResp.Body.Close() + + // Step 5: Remove the collection + deleteReq, err := http.NewRequest("DELETE", serverURL+baseURL+"/kb/collections/"+testCollectionID, nil) + assert.NoError(t, err) + deleteReq.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + deleteResp, err := http.DefaultClient.Do(deleteReq) + assert.NoError(t, err) + defer deleteResp.Body.Close() + + // Step 6: Verify collection no longer exists + finalExistsReq, err := http.NewRequest("GET", serverURL+baseURL+"/kb/collections/"+testCollectionID+"/exists", nil) + assert.NoError(t, err) + finalExistsReq.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken) + + finalExistsResp, err := http.DefaultClient.Do(finalExistsReq) + assert.NoError(t, err) + defer finalExistsResp.Body.Close() + + t.Logf("Completed full collection lifecycle test for: %s", testCollectionID) + }) +} diff --git a/openapi/tests/testutils/testutils.go b/openapi/tests/testutils/testutils.go index c99d2aab..5b8628c6 100644 --- a/openapi/tests/testutils/testutils.go +++ b/openapi/tests/testutils/testutils.go @@ -14,6 +14,7 @@ import ( "github.com/gin-gonic/gin" "github.com/yaoapp/yao/config" + "github.com/yaoapp/yao/kb" "github.com/yaoapp/yao/openapi" "github.com/yaoapp/yao/openapi/oauth/types" "github.com/yaoapp/yao/test" @@ -28,6 +29,52 @@ var testMutex sync.RWMutex // activeTestCount tracks the number of active tests using the global server var activeTestCount int +// testCollections tracks collections created during tests for cleanup +var testCollections []string +var testCollectionsMutex sync.Mutex + +// RegisterTestCollection adds a collection ID to the cleanup list +func RegisterTestCollection(collectionID string) { + testCollectionsMutex.Lock() + defer testCollectionsMutex.Unlock() + testCollections = append(testCollections, collectionID) +} + +// CleanupTestCollections removes all registered test collections +func CleanupTestCollections(t *testing.T) { + testCollectionsMutex.Lock() + collectionsToClean := make([]string, len(testCollections)) + copy(collectionsToClean, testCollections) + testCollections = nil // Clear the list + testCollectionsMutex.Unlock() + + if len(collectionsToClean) == 0 { + return + } + + // Check if KB instance is available + if kb.Instance == nil { + t.Logf("Warning: KB instance not available, cannot cleanup %d test collections", len(collectionsToClean)) + return + } + + ctx := context.Background() + cleanedCount := 0 + for _, collectionID := range collectionsToClean { + removed, err := kb.Instance.RemoveCollection(ctx, collectionID) + if err != nil { + t.Logf("Warning: Failed to cleanup test collection %s: %v", collectionID, err) + } else if removed { + cleanedCount++ + t.Logf("Cleaned up test collection: %s", collectionID) + } + } + + if cleanedCount > 0 { + t.Logf("Successfully cleaned up %d/%d test collections", cleanedCount, len(collectionsToClean)) + } +} + // Prepare initializes the OpenAPI test environment and starts a mock HTTP server. // // AI ASSISTANT INSTRUCTIONS: @@ -102,6 +149,14 @@ func Prepare(t *testing.T) string { // Step 1: Initialize base test environment with all Yao dependencies test.Prepare(t, config.Conf) + // Step 1.5: Initialize Knowledge Base (must be done before OpenAPI load) + _, err := kb.Load(config.Conf) + if err != nil { + // KB loading failure is not fatal for tests, just log it + t.Logf("Warning: Failed to load Knowledge Base: %v", err) + t.Logf("Some KB-related tests may not work properly") + } + // Step 2: Initialize OpenAPI server and make it available globally (only if not already initialized) if openapi.Server == nil { _, err := openapi.Load(config.Conf) @@ -163,7 +218,11 @@ func Prepare(t *testing.T) string { // // This ensures no state leakage between tests and prevents memory leaks // -// Step 3: Calls test.Clean() to clean up the base test environment +// Step 3: Clean up GraphRag test collections +// +// This removes any test collections that were created during the test +// +// Step 4: Calls test.Clean() to clean up the base test environment // // This closes database connections, cleans up temporary files, and resets global state // @@ -196,7 +255,29 @@ func Clean() { } testMutex.Unlock() - // Step 3: Clean up base test environment + // Step 3: Clean up GraphRag test collections (only when cleaning global state) + if shouldCleanGlobalState { + // Create a dummy test object for logging (since Clean doesn't receive *testing.T) + // Note: This is a limitation, but we can still attempt cleanup + ctx := context.Background() + testCollectionsMutex.Lock() + collectionsToClean := make([]string, len(testCollections)) + copy(collectionsToClean, testCollections) + testCollections = nil // Clear the list + testCollectionsMutex.Unlock() + + if len(collectionsToClean) > 0 && kb.Instance != nil { + for _, collectionID := range collectionsToClean { + _, err := kb.Instance.RemoveCollection(ctx, collectionID) + if err != nil { + // Can't use t.Logf here, but errors will be visible in test output + _ = err + } + } + } + } + + // Step 4: Clean up base test environment if shouldCleanGlobalState { test.Clean() }