Refactor Test Files and Enhance Test Environment Initialization

- Reorganized import statements in test files to improve clarity and consistency.
- Updated comments in test cases to provide more detailed descriptions of the test environment initialization process.
- Enhanced the `Prepare` function in the test utilities to include registration of the default query engine, ensuring proper setup for database searches.
- Improved error handling during the loading of the knowledge base and query engine, enhancing robustness in test setups.

agent/search/handlers/db/handler_integration_test.go

@@ -3,8 +3,6 @@ package db_test
 import (
 	"testing"
 
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
 	"github.com/yaoapp/gou/model"
 	"github.com/yaoapp/gou/query/gou"
 	"github.com/yaoapp/yao/agent/context"
@@ -12,6 +10,9 @@ import (
 	"github.com/yaoapp/yao/agent/search/types"
 	"github.com/yaoapp/yao/agent/testutils"
 	oauthTypes "github.com/yaoapp/yao/openapi/oauth/types"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 // ============================================================================
@@ -24,7 +25,7 @@ func TestHandler_Search_Integration(t *testing.T) {
 		t.Skip("Skipping integration test")
 	}
 
-	// Initialize test environment (loads models, database, etc.)
+	// Initialize test environment (loads models, database, query engine, etc.)
 	testutils.Prepare(t)
 	defer testutils.Clean(t)
 

agent/search/jsapi_db_test.go

@@ -3,14 +3,15 @@ package search_test
 import (
 	"testing"
 
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
 	"github.com/yaoapp/gou/model"
 	"github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/agent/search"
 	"github.com/yaoapp/yao/agent/search/types"
 	"github.com/yaoapp/yao/agent/testutils"
 	oauthTypes "github.com/yaoapp/yao/openapi/oauth/types"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 // ============================================================================
@@ -23,7 +24,7 @@ func TestJSAPI_DB_Integration(t *testing.T) {
 		t.Skip("Skipping integration test")
 	}
 
-	// Initialize test environment
+	// Initialize test environment (loads models, database, query engine, etc.)
 	testutils.Prepare(t)
 	defer testutils.Clean(t)
 

agent/testutils/testutils.go

@@ -4,11 +4,14 @@ import (
 	"testing"
 
 	_ "github.com/yaoapp/gou/encoding"
+	"github.com/yaoapp/gou/model"
+	"github.com/yaoapp/gou/query"
+	"github.com/yaoapp/gou/query/gou"
 	_ "github.com/yaoapp/gou/text"
+	"github.com/yaoapp/xun/capsule"
 	"github.com/yaoapp/yao/agent"
 	"github.com/yaoapp/yao/config"
 	"github.com/yaoapp/yao/kb"
-	"github.com/yaoapp/yao/query"
 	"github.com/yaoapp/yao/test"
 )
 
@@ -20,14 +23,8 @@ import (
 func Prepare(t *testing.T, opts ...interface{}) {
 	test.Prepare(t, config.Conf, opts...)
 
-	// Load Query Engine (required for DB search)
-	err := query.Load(config.Conf)
-	if err != nil {
-		t.Fatal(err)
-	}
-
 	// Load KB (required for agent KB features)
-	_, err = kb.Load(config.Conf)
+	_, err := kb.Load(config.Conf)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -37,6 +34,21 @@ func Prepare(t *testing.T, opts ...interface{}) {
 	if err != nil {
 		t.Fatal(err)
 	}
+
+	// Register default query engine (required for DB search)
+	// capsule.Global is initialized by test.Prepare
+	if _, has := query.Engines["default"]; !has && capsule.Global != nil {
+		query.Register("default", &gou.Query{
+			Query: capsule.Query(),
+			GetTableName: func(s string) string {
+				if mod, has := model.Models[s]; has {
+					return mod.MetaData.Table.Name
+				}
+				return s
+			},
+			AESKey: config.Conf.DB.AESKey,
+		})
+	}
 }
 
 // Clean clean the test environment