|
|
@@ -0,0 +1,121 @@ |
|
|
|
|
|
package model |
|
|
|
|
|
|
|
|
|
|
|
import ( |
|
|
|
|
|
"testing" |
|
|
|
|
|
|
|
|
|
|
|
"github.com/QuantumNous/new-api/common" |
|
|
|
|
|
"github.com/glebarez/sqlite" |
|
|
|
|
|
"github.com/stretchr/testify/assert" |
|
|
|
|
|
"github.com/stretchr/testify/require" |
|
|
|
|
|
"gorm.io/gorm" |
|
|
|
|
|
) |
|
|
|
|
|
|
|
|
|
|
|
func setupChannelPricingDB(t *testing.T) *gorm.DB { |
|
|
|
|
|
t.Helper() |
|
|
|
|
|
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) |
|
|
|
|
|
require.NoError(t, err) |
|
|
|
|
|
sqlDB, _ := db.DB() |
|
|
|
|
|
sqlDB.SetMaxOpenConns(1) |
|
|
|
|
|
|
|
|
|
|
|
origDB := DB |
|
|
|
|
|
DB = db |
|
|
|
|
|
common.UsingSQLite = true |
|
|
|
|
|
common.RedisEnabled = false |
|
|
|
|
|
|
|
|
|
|
|
require.NoError(t, db.AutoMigrate(&ChannelPricing{})) |
|
|
|
|
|
|
|
|
|
|
|
t.Cleanup(func() { |
|
|
|
|
|
DB = origDB |
|
|
|
|
|
sqlDB.Close() |
|
|
|
|
|
}) |
|
|
|
|
|
return db |
|
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
func TestCacheWriteThrough_Insert(t *testing.T) { |
|
|
|
|
|
setupChannelPricingDB(t) |
|
|
|
|
|
|
|
|
|
|
|
cp := &ChannelPricing{ |
|
|
|
|
|
ModelName: "test-insert-model", |
|
|
|
|
|
ChannelId: 9001, |
|
|
|
|
|
QuotaType: QuotaTypeByTokens, |
|
|
|
|
|
ModelRatio: 1.0, |
|
|
|
|
|
CacheRatio: 0.8, |
|
|
|
|
|
ImageRatio: 1.2, |
|
|
|
|
|
} |
|
|
|
|
|
require.NoError(t, cp.Insert()) |
|
|
|
|
|
defer cp.Delete() |
|
|
|
|
|
|
|
|
|
|
|
found, ok := GetEffectivePricing("test-insert-model", 9001) |
|
|
|
|
|
assert.True(t, ok) |
|
|
|
|
|
assert.Equal(t, 0.8, found.CacheRatio) |
|
|
|
|
|
assert.Equal(t, 1.2, found.ImageRatio) |
|
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
func TestCacheWriteThrough_Update(t *testing.T) { |
|
|
|
|
|
setupChannelPricingDB(t) |
|
|
|
|
|
|
|
|
|
|
|
cp := &ChannelPricing{ |
|
|
|
|
|
ModelName: "test-update-model", |
|
|
|
|
|
ChannelId: 9002, |
|
|
|
|
|
QuotaType: QuotaTypeByTokens, |
|
|
|
|
|
ModelRatio: 1.0, |
|
|
|
|
|
} |
|
|
|
|
|
require.NoError(t, cp.Insert()) |
|
|
|
|
|
defer cp.Delete() |
|
|
|
|
|
|
|
|
|
|
|
cp.CacheRatio = 0.9 |
|
|
|
|
|
cp.AudioRatio = 1.5 |
|
|
|
|
|
require.NoError(t, cp.Update()) |
|
|
|
|
|
|
|
|
|
|
|
found, ok := GetEffectivePricing("test-update-model", 9002) |
|
|
|
|
|
assert.True(t, ok) |
|
|
|
|
|
assert.Equal(t, 0.9, found.CacheRatio) |
|
|
|
|
|
assert.Equal(t, 1.5, found.AudioRatio) |
|
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
func TestCacheWriteThrough_Delete(t *testing.T) { |
|
|
|
|
|
setupChannelPricingDB(t) |
|
|
|
|
|
|
|
|
|
|
|
cp := &ChannelPricing{ |
|
|
|
|
|
ModelName: "test-delete-model", |
|
|
|
|
|
ChannelId: 9003, |
|
|
|
|
|
QuotaType: QuotaTypeByTokens, |
|
|
|
|
|
ModelRatio: 1.0, |
|
|
|
|
|
} |
|
|
|
|
|
require.NoError(t, cp.Insert()) |
|
|
|
|
|
|
|
|
|
|
|
require.NoError(t, cp.Delete()) |
|
|
|
|
|
|
|
|
|
|
|
found, ok := GetEffectivePricing("test-delete-model", 9003) |
|
|
|
|
|
assert.False(t, ok) |
|
|
|
|
|
assert.Nil(t, found) |
|
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
func TestExtendedFields_DefaultZero(t *testing.T) { |
|
|
|
|
|
setupChannelPricingDB(t) |
|
|
|
|
|
|
|
|
|
|
|
cp := &ChannelPricing{ |
|
|
|
|
|
ModelName: "test-default-model", |
|
|
|
|
|
ChannelId: 9004, |
|
|
|
|
|
QuotaType: QuotaTypeByTokens, |
|
|
|
|
|
ModelRatio: 1.0, |
|
|
|
|
|
} |
|
|
|
|
|
require.NoError(t, cp.Insert()) |
|
|
|
|
|
defer cp.Delete() |
|
|
|
|
|
|
|
|
|
|
|
found, ok := GetEffectivePricing("test-default-model", 9004) |
|
|
|
|
|
assert.True(t, ok) |
|
|
|
|
|
assert.Equal(t, 0.0, found.CacheRatio) |
|
|
|
|
|
assert.Equal(t, 0.0, found.CacheCreationRatio) |
|
|
|
|
|
assert.Equal(t, 0.0, found.ImageRatio) |
|
|
|
|
|
assert.Equal(t, 0.0, found.AudioRatio) |
|
|
|
|
|
assert.Equal(t, 0.0, found.AudioCompletionRatio) |
|
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
func TestGetEffectivePricing_NotFound(t *testing.T) { |
|
|
|
|
|
setupChannelPricingDB(t) |
|
|
|
|
|
|
|
|
|
|
|
found, ok := GetEffectivePricing("nonexistent-model-xyz", 99999) |
|
|
|
|
|
assert.False(t, ok) |
|
|
|
|
|
assert.Nil(t, found) |
|
|
|
|
|
} |