diff --git a/agent/search/handlers/db/handler_integration_test.go b/agent/search/handlers/db/handler_integration_test.go index 5d6bf2e5..9ccfa5a3 100644 --- a/agent/search/handlers/db/handler_integration_test.go +++ b/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) diff --git a/agent/search/jsapi_db_test.go b/agent/search/jsapi_db_test.go index 5a8160a1..5c315e68 100644 --- a/agent/search/jsapi_db_test.go +++ b/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) diff --git a/agent/testutils/testutils.go b/agent/testutils/testutils.go index ef9d246b..104ad06a 100644 --- a/agent/testutils/testutils.go +++ b/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