mirror of
https://github.com/Rain-kl/OpenFlare.git
synced 2026-09-29 22:06:38 +08:00
ryan
399 lines
15 KiB
Go
399 lines
15 KiB
Go
// Copyright 2025 linux.do
|
|
// Copyright 2026 Arctel.net
|
|
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
// Package testhelper 提供测试辅助工具
|
|
package testhelper
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"io"
|
|
"strconv"
|
|
"testing"
|
|
"time"
|
|
|
|
"Wavelet/core/contracts"
|
|
"Wavelet/openflare/plugins/server/kernel/model"
|
|
analyticsmodel "Wavelet/openflare/plugins/server/kernel/model/analytics"
|
|
"Wavelet/openflare/plugins/server/kernel/ofupload"
|
|
"Wavelet/openflare/plugins/server/kernel/repository"
|
|
"Wavelet/openflare/plugins/server/kernel/repository/logstore"
|
|
oftask "Wavelet/openflare/plugins/server/kernel/task"
|
|
"Wavelet/pkg/idgen"
|
|
|
|
"github.com/glebarez/sqlite"
|
|
"github.com/stretchr/testify/require"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
const (
|
|
configTypeSystem = "system"
|
|
configTypeBusiness = "business"
|
|
configValueTrue = "true"
|
|
configValueFalse = "false"
|
|
)
|
|
|
|
type testConfigService struct {
|
|
db *gorm.DB
|
|
}
|
|
|
|
// NewMockSystemConfigService creates a test SystemConfigService backed by GORM.
|
|
func NewMockSystemConfigService(db *gorm.DB) contracts.SystemConfigService {
|
|
return &testConfigService{db: db}
|
|
}
|
|
|
|
func (s *testConfigService) GetByKey(ctx context.Context, key string) (contracts.SystemConfigDTO, error) {
|
|
var cfg contracts.SystemConfigDTO
|
|
err := s.db.WithContext(ctx).Table("w_system_configs").Where("key = ?", key).First(&cfg).Error
|
|
return cfg, err
|
|
}
|
|
|
|
func (s *testConfigService) ListByKeys(ctx context.Context, keys []string) (map[string]contracts.SystemConfigDTO, error) {
|
|
var cfgs []contracts.SystemConfigDTO
|
|
if err := s.db.WithContext(ctx).Table("w_system_configs").Where("key IN ?", keys).Find(&cfgs).Error; err != nil {
|
|
return nil, err
|
|
}
|
|
res := make(map[string]contracts.SystemConfigDTO, len(cfgs))
|
|
for _, c := range cfgs {
|
|
res[c.Key] = c
|
|
}
|
|
return res, nil
|
|
}
|
|
|
|
func (s *testConfigService) ListVisible(ctx context.Context) ([]contracts.SystemConfigDTO, error) {
|
|
var cfgs []contracts.SystemConfigDTO
|
|
err := s.db.WithContext(ctx).Table("w_system_configs").Where("visibility = ?", 1).Find(&cfgs).Error
|
|
return cfgs, err
|
|
}
|
|
|
|
func (s *testConfigService) ListByType(ctx context.Context, configType string) ([]contracts.SystemConfigDTO, error) {
|
|
var cfgs []contracts.SystemConfigDTO
|
|
err := s.db.WithContext(ctx).Table("w_system_configs").Where("type = ?", configType).Find(&cfgs).Error
|
|
return cfgs, err
|
|
}
|
|
|
|
func (s *testConfigService) GetIntByKey(ctx context.Context, key string) (int, error) {
|
|
var cfg contracts.SystemConfigDTO
|
|
if err := s.db.WithContext(ctx).Table("w_system_configs").Where("key = ?", key).First(&cfg).Error; err != nil {
|
|
return 0, err
|
|
}
|
|
return strconv.Atoi(cfg.Value)
|
|
}
|
|
|
|
func (s *testConfigService) GetBoolByKey(ctx context.Context, key string) (bool, error) {
|
|
var cfg contracts.SystemConfigDTO
|
|
if err := s.db.WithContext(ctx).Table("w_system_configs").Where("key = ?", key).First(&cfg).Error; err != nil {
|
|
return false, err
|
|
}
|
|
return strconv.ParseBool(cfg.Value)
|
|
}
|
|
|
|
func (s *testConfigService) SaveOrUpdate(ctx context.Context, key, value string) error {
|
|
var cfg model.SystemConfig
|
|
if err := s.db.WithContext(ctx).Table("w_system_configs").Where("key = ?", key).First(&cfg).Error; err != nil {
|
|
cfg = model.SystemConfig{Key: key, Value: value, Type: "system"}
|
|
return s.db.WithContext(ctx).Table("w_system_configs").Create(&cfg).Error
|
|
}
|
|
return s.db.WithContext(ctx).Table("w_system_configs").Where("key = ?", key).Update("value", value).Error
|
|
}
|
|
|
|
func (s *testConfigService) InvalidateCache(ctx context.Context, key string) error { return nil }
|
|
func (s *testConfigService) InvalidateAllCaches(ctx context.Context) error { return nil }
|
|
|
|
type testSystemConfigEntity struct {
|
|
Key string `gorm:"primaryKey"`
|
|
Value string
|
|
Type string
|
|
Visibility int
|
|
Description string
|
|
UpdatedAt time.Time
|
|
CreatedAt time.Time
|
|
}
|
|
|
|
func (testSystemConfigEntity) TableName() string {
|
|
return "w_system_configs"
|
|
}
|
|
|
|
// SetupTestEnvironment initializes an in-memory SQLite DB and seeds default
|
|
// configurations. Redis is no longer owned by OpenFlare.
|
|
func SetupTestEnvironment(t *testing.T) (*gorm.DB, any, func()) {
|
|
t.Helper()
|
|
sqliteDB, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{
|
|
DisableForeignKeyConstraintWhenMigrating: true,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("failed to open in-memory SQLite db: %v", err)
|
|
}
|
|
|
|
if sqlDB, err := sqliteDB.DB(); err == nil {
|
|
sqlDB.SetMaxOpenConns(1)
|
|
}
|
|
|
|
err = sqliteDB.AutoMigrate(
|
|
&testSystemConfigEntity{},
|
|
&model.User{},
|
|
&model.AccessToken{},
|
|
&model.Upload{},
|
|
&model.UploadStat{},
|
|
&model.TaskExecution{},
|
|
)
|
|
if err != nil {
|
|
t.Fatalf("failed to auto migrate tables: %v", err)
|
|
}
|
|
|
|
repository.SetDBForTest(sqliteDB)
|
|
repository.SetSystemConfigService(&testConfigService{db: sqliteDB})
|
|
|
|
mockStorage := NewMockStorageService()
|
|
ofupload.SetStorage(mockStorage)
|
|
ofupload.SetUploadService(&mockUploadService{db: sqliteDB})
|
|
noopTask := &NoopTaskService{}
|
|
repository.SetTaskService(noopTask)
|
|
oftask.SetService(noopTask)
|
|
|
|
if err := idgen.Init(1); err != nil {
|
|
t.Fatalf("idgen.Init: %v", err)
|
|
}
|
|
seedDefaultConfigs(t, sqliteDB)
|
|
|
|
cleanup := func() {
|
|
runExtraCleanups()
|
|
repository.StopSystemConfigCacheListener()
|
|
repository.SetAuthService(nil)
|
|
repository.SetUserService(nil)
|
|
repository.SetSystemConfigService(nil)
|
|
repository.SetTaskService(nil)
|
|
repository.SetDBForTest(nil)
|
|
ofupload.SetStorage(nil)
|
|
ofupload.SetUploadService(nil)
|
|
oftask.SetService(nil)
|
|
}
|
|
|
|
return sqliteDB, nil, cleanup
|
|
}
|
|
|
|
type mockUploadService struct {
|
|
db *gorm.DB
|
|
}
|
|
|
|
// NewMockUploadService creates a mock UploadService backed by GORM.
|
|
func NewMockUploadService(db *gorm.DB) contracts.UploadService {
|
|
return &mockUploadService{db: db}
|
|
}
|
|
|
|
func (s *mockUploadService) GetByID(ctx context.Context, id uint64) (*contracts.UploadDTO, error) {
|
|
var u contracts.UploadDTO
|
|
err := s.db.WithContext(ctx).Table("w_uploads").Where("id = ?", id).First(&u).Error
|
|
if err != nil {
|
|
return &contracts.UploadDTO{
|
|
ID: id,
|
|
Status: "used",
|
|
Type: "openflare_pages_deployment",
|
|
Size: 100,
|
|
CreatedAt: time.Now().UTC(),
|
|
UpdatedAt: time.Now().UTC(),
|
|
}, nil
|
|
}
|
|
return &u, nil
|
|
}
|
|
|
|
func (s *mockUploadService) OpenStoredUpload(ctx context.Context, id uint64) (*contracts.OpenedUploadDTO, error) {
|
|
u, err := s.GetByID(ctx, id)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
body := io.ReadCloser(io.NopCloser(bytes.NewReader(nil)))
|
|
storage := ofupload.CurrentStorage()
|
|
if storage != nil {
|
|
if obj, err := storage.Get(ctx, u.FilePath); err == nil && obj != nil && obj.Body != nil {
|
|
body = obj.Body
|
|
} else if obj, err := storage.Get(ctx, u.FileName); err == nil && obj != nil && obj.Body != nil {
|
|
body = obj.Body
|
|
}
|
|
}
|
|
return &contracts.OpenedUploadDTO{
|
|
Upload: *u,
|
|
Body: body,
|
|
ContentType: u.MimeType,
|
|
ContentLength: u.Size,
|
|
}, nil
|
|
}
|
|
|
|
func (s *mockUploadService) Remove(ctx context.Context, id uint64) error {
|
|
if err := s.db.WithContext(ctx).Table("w_uploads").Where("id = ?", id).Update("status", "deleted").Error; err != nil {
|
|
return err
|
|
}
|
|
return s.RebuildStats(ctx)
|
|
}
|
|
|
|
func (s *mockUploadService) RemoveOwned(ctx context.Context, id uint64, userID uint64) error {
|
|
if err := s.db.WithContext(ctx).Table("w_uploads").Where("id = ? AND user_id = ?", id, userID).Update("status", "deleted").Error; err != nil {
|
|
return err
|
|
}
|
|
return s.RebuildStats(ctx)
|
|
}
|
|
|
|
func (s *mockUploadService) FindByHash(ctx context.Context, hash string, size int64) (*contracts.UploadDTO, error) {
|
|
var u contracts.UploadDTO
|
|
err := s.db.WithContext(ctx).Table("w_uploads").Where("hash = ? AND size = ?", hash, size).First(&u).Error
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &u, nil
|
|
}
|
|
|
|
func (s *mockUploadService) RebuildStats(ctx context.Context) error {
|
|
var count int64
|
|
_ = s.db.WithContext(ctx).Table("w_uploads").Where("status != ?", "deleted").Count(&count).Error
|
|
var stat model.UploadStat
|
|
if err := s.db.WithContext(ctx).Table("w_upload_stats").Where("dimension = ?", model.UploadStatDimensionTotal).First(&stat).Error; err != nil {
|
|
stat = model.UploadStat{
|
|
Dimension: model.UploadStatDimensionTotal,
|
|
FileCount: int(count),
|
|
}
|
|
return s.db.WithContext(ctx).Table("w_upload_stats").Create(&stat).Error
|
|
}
|
|
stat.FileCount = int(count)
|
|
return s.db.WithContext(ctx).Table("w_upload_stats").Save(&stat).Error
|
|
}
|
|
|
|
func getSeedConfigsPart1() []model.SystemConfig {
|
|
return []model.SystemConfig{
|
|
{Key: model.ConfigKeyUploadAllowedExtensions, Value: "jpg,png,webp", Type: configTypeSystem, Description: "允许上传的图片扩展名（逗号分隔）"},
|
|
{Key: model.ConfigKeySiteName, Value: "OpenFlare", Type: configTypeSystem, Description: "系统平台的展示名称"},
|
|
{Key: model.ConfigKeyPasswordLoginEnabled, Value: configValueTrue, Type: configTypeSystem},
|
|
{Key: model.ConfigKeyRegistrationEnabled, Value: configValueFalse, Type: configTypeSystem},
|
|
{Key: model.ConfigKeyPasswordRegisterEnabled, Value: configValueFalse, Type: configTypeSystem},
|
|
{Key: model.ConfigKeyOIDCLoginEnabled, Value: configValueTrue, Type: configTypeSystem},
|
|
{Key: model.ConfigKeyMaxAPIKeysPerUser, Value: "5", Type: "business"},
|
|
{Key: model.ConfigKeyCapLoginEnabled, Value: configValueFalse, Type: configTypeSystem},
|
|
{Key: model.ConfigKeyCapAutoSolve, Value: configValueTrue, Type: configTypeSystem},
|
|
{Key: model.ConfigKeyCapChallengeCount, Value: "1", Type: configTypeSystem},
|
|
{Key: model.ConfigKeyCapChallengeSize, Value: "32", Type: configTypeSystem},
|
|
{Key: model.ConfigKeyCapChallengeDifficulty, Value: "4", Type: configTypeSystem},
|
|
{Key: model.ConfigKeyCapChallengeTTL, Value: "600", Type: configTypeSystem},
|
|
{Key: model.ConfigKeyCapTokenTTL, Value: "1200", Type: configTypeSystem},
|
|
}
|
|
}
|
|
|
|
func getSeedConfigsPart2() []model.SystemConfig {
|
|
return []model.SystemConfig{
|
|
{Key: model.ConfigKeyServerAddress, Value: "", Type: configTypeSystem},
|
|
{Key: model.ConfigKeySMTPHost, Value: "", Type: configTypeSystem},
|
|
{Key: model.ConfigKeySMTPPort, Value: "587", Type: configTypeSystem},
|
|
{Key: model.ConfigKeySMTPUsername, Value: "", Type: configTypeSystem},
|
|
{Key: model.ConfigKeySMTPPassword, Value: "", Type: configTypeSystem},
|
|
{Key: model.ConfigKeyEmailLoginVerificationEnabled, Value: configValueFalse, Type: configTypeSystem},
|
|
{Key: model.ConfigKeyEmailRegisterVerificationEnabled, Value: configValueFalse, Type: configTypeSystem},
|
|
{Key: model.ConfigKeyMenuDisplayConfig, Value: "{}", Type: configTypeSystem},
|
|
{Key: model.ConfigKeySearchEngineIndexingEnabled, Value: configValueFalse, Type: configTypeSystem},
|
|
{Key: model.ConfigKeyFileAccessWhitelist, Value: `["avatar"]`, Type: configTypeSystem},
|
|
{Key: model.ConfigKeyDiskCacheMaxSizeMB, Value: "100", Type: configTypeSystem},
|
|
{Key: model.ConfigKeyDiskCacheTTLMinutes, Value: "60", Type: configTypeSystem},
|
|
{Key: model.ConfigKeyDiskCacheLRUEnabled, Value: configValueTrue, Type: configTypeSystem},
|
|
{Key: model.ConfigKeyLoginSessionTTLHours, Value: "0", Type: configTypeSystem},
|
|
{Key: model.ConfigKeyUpdateUpstreamRepository, Value: "Rain-kl/OpenFlare", Type: configTypeSystem},
|
|
{Key: model.ConfigKeyStorageConfig, Value: `{"driver":"local","local":{"root":"."},"s3":{"region":"us-east-1"},"r2":{"region":"auto"},"minio":{"region":"us-east-1","path_style":true},"oss":{},"webdav":{}}`, Type: configTypeSystem},
|
|
{Key: model.ConfigKeyRelayFRPSWebUIEnabled, Value: configValueFalse, Type: configTypeBusiness},
|
|
{Key: model.ConfigKeyRelayFRPSWebUIPort, Value: "17500", Type: configTypeBusiness},
|
|
{Key: model.ConfigKeyPagesMaxPackageSizeMB, Value: "100", Type: configTypeBusiness},
|
|
{Key: model.ConfigKeyPagesMaxHistoryCount, Value: "20", Type: configTypeBusiness},
|
|
{Key: model.ConfigKeyLogRetentionDaysPostgres, Value: "90", Type: configTypeBusiness},
|
|
{Key: model.ConfigKeyLogRetentionDaysSQLite, Value: "90", Type: configTypeBusiness},
|
|
{Key: model.ConfigKeyLogRetentionDaysClickHouse, Value: "90", Type: configTypeBusiness},
|
|
}
|
|
}
|
|
|
|
func seedDefaultConfigs(t *testing.T, tx *gorm.DB) {
|
|
t.Helper()
|
|
defaultConfigs := append(getSeedConfigsPart1(), getSeedConfigsPart2()...)
|
|
if err := tx.Table("w_system_configs").Create(&defaultConfigs).Error; err != nil {
|
|
t.Fatalf("failed to seed default system configs: %v", err)
|
|
}
|
|
|
|
publicKeys := []string{
|
|
model.ConfigKeyUploadAllowedExtensions,
|
|
model.ConfigKeySiteName,
|
|
model.ConfigKeyPasswordLoginEnabled,
|
|
model.ConfigKeyRegistrationEnabled,
|
|
model.ConfigKeyPasswordRegisterEnabled,
|
|
model.ConfigKeyOIDCLoginEnabled,
|
|
model.ConfigKeyMaxAPIKeysPerUser,
|
|
model.ConfigKeyCapLoginEnabled,
|
|
model.ConfigKeyCapAutoSolve,
|
|
model.ConfigKeyEmailLoginVerificationEnabled,
|
|
model.ConfigKeyEmailRegisterVerificationEnabled,
|
|
model.ConfigKeyMenuDisplayConfig,
|
|
model.ConfigKeySearchEngineIndexingEnabled,
|
|
model.ConfigKeyFileAccessWhitelist,
|
|
}
|
|
if err := tx.Table("w_system_configs").
|
|
Where("key IN ?", publicKeys).
|
|
Update("visibility", model.ConfigVisibilityVisible).Error; err != nil {
|
|
t.Fatalf("failed to seed public system config visibility: %v", err)
|
|
}
|
|
}
|
|
|
|
// SetupLogStoresForTest 将 logstore 指向测试已通过 SetDBForTest 注入的 sqlite 库。
|
|
func SetupLogStoresForTest(t *testing.T) {
|
|
t.Helper()
|
|
|
|
gdb := repository.DB(context.Background())
|
|
require.NoError(t, idgen.Init(1))
|
|
require.NoError(t, gdb.AutoMigrate(
|
|
&analyticsmodel.NodeAccessLog{},
|
|
&analyticsmodel.UserAccessLog{},
|
|
&analyticsmodel.NodeMetricSnapshot{},
|
|
&analyticsmodel.NodeEdgeHealth{},
|
|
&analyticsmodel.NodeObsFrps{},
|
|
&analyticsmodel.NodeObsFrpc{},
|
|
))
|
|
|
|
logstore.ResetForTest()
|
|
logstore.SetConfigReader(func(_ context.Context, key string) (string, error) {
|
|
if key == model.ConfigKeyLogDatabase {
|
|
return "sqlite", nil
|
|
}
|
|
return "", nil
|
|
})
|
|
store, err := logstore.Active(context.Background())
|
|
require.NoError(t, err)
|
|
|
|
logstore.SetAccessLogHooks(logstore.AccessLogHooks{
|
|
QueueNodeAccessLogs: func(logs []analyticsmodel.NodeAccessLog) {
|
|
if err := store.AccessLogs.BatchInsertNodeAccessLogs(context.Background(), logs); err != nil {
|
|
t.Errorf("batch insert node access logs failed in test hook: %v", err)
|
|
}
|
|
},
|
|
})
|
|
logstore.SetObservabilityHooks(logstore.ObservabilityHooks{
|
|
QueueMetricSnapshot: func(record analyticsmodel.NodeMetricSnapshot) {
|
|
if err := store.Observability.BatchInsertNodeMetricSnapshots(context.Background(), []analyticsmodel.NodeMetricSnapshot{record}); err != nil {
|
|
t.Errorf("batch insert node metric snapshots failed in test hook: %v", err)
|
|
}
|
|
},
|
|
QueueEdgeHealth: func(record analyticsmodel.NodeEdgeHealth) {
|
|
if err := store.Observability.BatchInsertNodeEdgeHealth(context.Background(), []analyticsmodel.NodeEdgeHealth{record}); err != nil {
|
|
t.Errorf("batch insert node edge health failed in test hook: %v", err)
|
|
}
|
|
},
|
|
QueueNodeObsFrps: func(record analyticsmodel.NodeObsFrps) {
|
|
if err := store.Observability.BatchInsertNodeObsFrps(context.Background(), []analyticsmodel.NodeObsFrps{record}); err != nil {
|
|
t.Errorf("batch insert node obs frps failed in test hook: %v", err)
|
|
}
|
|
},
|
|
QueueNodeObsFrpc: func(record analyticsmodel.NodeObsFrpc) {
|
|
if err := store.Observability.BatchInsertNodeObsFrpc(context.Background(), []analyticsmodel.NodeObsFrpc{record}); err != nil {
|
|
t.Errorf("batch insert node obs frpc failed in test hook: %v", err)
|
|
}
|
|
},
|
|
})
|
|
|
|
t.Cleanup(func() {
|
|
logstore.SetAccessLogHooks(logstore.AccessLogHooks{})
|
|
logstore.SetObservabilityHooks(logstore.ObservabilityHooks{})
|
|
logstore.ResetForTest()
|
|
})
|
|
}
|