// // Copyright 2026 The InfiniFlow Authors. All Rights Reserved. // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. // package task import ( "fmt" "strings" "testing" "ragflow/internal/dao" "ragflow/internal/entity" "github.com/glebarez/sqlite" "gorm.io/gorm" ) func setupTestDB(t *testing.T) *gorm.DB { t.Helper() dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.NewReplacer("/", "_", " ", "_").Replace(t.Name())) db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{TranslateError: true}) if err != nil { t.Fatalf("open sqlite: %v", err) } sqlDB, err := db.DB() if err != nil { t.Fatalf("failed to get sql DB: %v", err) } sqlDB.SetMaxOpenConns(1) if err := db.AutoMigrate( &entity.Task{}, &entity.Document{}, &entity.Knowledgebase{}, &entity.Tenant{}, ); err != nil { t.Fatalf("migrate: %v", err) } return db } func pushDB(t *testing.T, db *gorm.DB) { t.Helper() orig := dao.DB dao.DB = db t.Cleanup(func() { dao.DB = orig }) } func ptr(s string) *string { return &s }