69 lines
1.9 KiB
Go
69 lines
1.9 KiB
Go
package controller
|
|
|
|
import (
|
|
"testing"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/model"
|
|
"github.com/glebarez/sqlite"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
func setupBalanceAlertTestDB(t *testing.T) {
|
|
t.Helper()
|
|
db, err := gorm.Open(sqlite.Open("file::memory:?cache=shared"), &gorm.Config{})
|
|
if err != nil {
|
|
t.Fatalf("open sqlite failed: %v", err)
|
|
}
|
|
model.DB = db
|
|
if err := db.AutoMigrate(&model.Channel{}, &model.UserMessage{}); err != nil {
|
|
t.Fatalf("auto migrate failed: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestNotifyChannelBalanceAlertIfNeeded(t *testing.T) {
|
|
setupBalanceAlertTestDB(t)
|
|
|
|
common.OptionMapRWMutex.Lock()
|
|
common.OptionMap = map[string]string{
|
|
"ChannelBalanceAlertEnabled": "true",
|
|
"ChannelBalanceSoftAlertThreshold": "50",
|
|
"ChannelBalanceRiskAlertThreshold": "20",
|
|
}
|
|
common.OptionMapRWMutex.Unlock()
|
|
|
|
ch := &model.Channel{
|
|
Name: "test-channel",
|
|
Balance: 100,
|
|
}
|
|
if err := model.DB.Create(ch).Error; err != nil {
|
|
t.Fatalf("create channel failed: %v", err)
|
|
}
|
|
|
|
notifyChannelBalanceAlertIfNeeded(ch, 100, 10)
|
|
|
|
var firstCount int64
|
|
if err := model.DB.Model(&model.UserMessage{}).Count(&firstCount).Error; err != nil {
|
|
t.Fatalf("count messages failed: %v", err)
|
|
}
|
|
if firstCount != 1 {
|
|
t.Fatalf("expected 1 message after entering risk level, got %d", firstCount)
|
|
}
|
|
var firstMsg model.UserMessage
|
|
if err := model.DB.Order("id desc").First(&firstMsg).Error; err != nil {
|
|
t.Fatalf("load latest message failed: %v", err)
|
|
}
|
|
if firstMsg.ReceiverMinRole != common.RoleAdminUser {
|
|
t.Fatalf("expected receiver_min_role=%d, got %d", common.RoleAdminUser, firstMsg.ReceiverMinRole)
|
|
}
|
|
|
|
notifyChannelBalanceAlertIfNeeded(ch, 10, 5)
|
|
var secondCount int64
|
|
if err := model.DB.Model(&model.UserMessage{}).Count(&secondCount).Error; err != nil {
|
|
t.Fatalf("count messages failed: %v", err)
|
|
}
|
|
if secondCount != 1 {
|
|
t.Fatalf("expected no duplicate message in same level, got %d", secondCount)
|
|
}
|
|
}
|