|
- 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)
- }
|