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