yao/agent/testutils/testutils.go
Max 19153ee4a5 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.
2025-12-18 12:11:51 +08:00

57 lines
1.4 KiB
Go

package testutils
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/test"
)
// Prepare prepare the test environment with optional V8 mode configuration
// Usage:
//
// testutils.Prepare(t) // standard mode (default)
// testutils.Prepare(t, test.PrepareOption{V8Mode: "performance"}) // performance mode for benchmarks
func Prepare(t *testing.T, opts ...interface{}) {
test.Prepare(t, config.Conf, opts...)
// Load KB (required for agent KB features)
_, err := kb.Load(config.Conf)
if err != nil {
t.Fatal(err)
}
// Load agent
err = agent.Load(config.Conf)
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
func Clean(t *testing.T) {
test.Clean()
}