93 lines
3.8 KiB
Go
93 lines
3.8 KiB
Go
package model
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"os"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/glebarez/sqlite"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"gorm.io/driver/mysql"
|
|
"gorm.io/driver/postgres"
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/schema"
|
|
)
|
|
|
|
func TestGroupCacheStatsDatabaseMatrix(t *testing.T) {
|
|
for _, dialect := range []string{"sqlite", "mysql", "postgres"} {
|
|
t.Run(dialect, func(t *testing.T) {
|
|
var driver gorm.Dialector
|
|
switch dialect {
|
|
case "sqlite":
|
|
driver = sqlite.Open(":memory:")
|
|
case "mysql":
|
|
dsn := os.Getenv("TEST_MYSQL_DSN")
|
|
if dsn == "" {
|
|
t.Skip("TEST_MYSQL_DSN is not configured")
|
|
}
|
|
driver = mysql.Open(dsn)
|
|
case "postgres":
|
|
dsn := os.Getenv("TEST_POSTGRES_DSN")
|
|
if dsn == "" {
|
|
t.Skip("TEST_POSTGRES_DSN is not configured")
|
|
}
|
|
driver = postgres.Open(dsn)
|
|
}
|
|
db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("gs_%d_", time.Now().UnixNano())}})
|
|
require.NoError(t, err)
|
|
originalDB, originalLogDB := DB, LOG_DB
|
|
DB, LOG_DB = db, db
|
|
t.Cleanup(func() { DB, LOG_DB = originalDB, originalLogDB })
|
|
channelTable := db.NamingStrategy.TableName("channels")
|
|
logTable := db.NamingStrategy.TableName("logs")
|
|
type sample struct {
|
|
Id int
|
|
ChannelId int
|
|
PromptTokens int64
|
|
Other string
|
|
Type int
|
|
TokenId int
|
|
CreatedAt int64
|
|
}
|
|
require.NoError(t, db.Table(channelTable).AutoMigrate(&StatusChannel{}))
|
|
require.NoError(t, db.Table(logTable).AutoMigrate(&sample{}))
|
|
t.Cleanup(func() {
|
|
require.NoError(t, db.Migrator().DropTable(channelTable, logTable))
|
|
sqlDB, err := db.DB()
|
|
require.NoError(t, err)
|
|
require.NoError(t, sqlDB.Close())
|
|
})
|
|
require.NoError(t, db.Table(channelTable).Create(&[]StatusChannel{
|
|
{Id: 1, Group: "renamed", Status: 1}, {Id: 2, Group: "automatic-down", Status: 3},
|
|
{Id: 3, Group: "manual-off", Status: 2}, {Id: 4, Group: "unused", Status: 1},
|
|
}).Error)
|
|
require.NoError(t, db.Table(logTable).Create(&[]sample{
|
|
{ChannelId: 1, PromptTokens: 100, Other: `{"cache_tokens":80}`, Type: 2, TokenId: 1, CreatedAt: 150},
|
|
{ChannelId: 1, PromptTokens: 300, Other: `{"cache_tokens":0}`, Type: 2, TokenId: 1, CreatedAt: 150},
|
|
{ChannelId: 1, PromptTokens: 20, Other: `{"cache_tokens":50,"usage_semantic":"anthropic","cache_creation_tokens":30}`, Type: 2, TokenId: 1, CreatedAt: 150},
|
|
{ChannelId: 1, PromptTokens: 10, Other: `{"cache_tokens":40,"input_tokens_total":100,"claude":true}`, Type: 2, TokenId: 1, CreatedAt: 150},
|
|
{ChannelId: 2, PromptTokens: 50, Other: `{}`, Type: 2, TokenId: 1, CreatedAt: 150},
|
|
{ChannelId: 1, PromptTokens: 999, Other: `{}`, Type: 2, TokenId: 0, CreatedAt: 150},
|
|
{ChannelId: 1, PromptTokens: 999, Other: `{}`, Type: 5, TokenId: 1, CreatedAt: 150},
|
|
{ChannelId: 1, PromptTokens: 999, Other: `{}`, Type: 2, TokenId: 1, CreatedAt: 99},
|
|
{ChannelId: 1, PromptTokens: 999, Other: `{}`, Type: 2, TokenId: 1, CreatedAt: 200},
|
|
{ChannelId: 1, PromptTokens: 999, Other: `invalid`, Type: 2, TokenId: 1, CreatedAt: 150},
|
|
{ChannelId: 1, PromptTokens: 0, Other: `{}`, Type: 2, TokenId: 1, CreatedAt: 150},
|
|
{ChannelId: 3, PromptTokens: 999, Other: `{}`, Type: 2, TokenId: 1, CreatedAt: 150},
|
|
}).Error)
|
|
channels, err := GetStatusChannels(context.Background())
|
|
require.NoError(t, err)
|
|
require.Len(t, channels, 3)
|
|
stats, err := GetGroupCacheStats(context.Background(), channels, time.Unix(100, 0), time.Unix(200, 0))
|
|
require.NoError(t, err)
|
|
assert.Equal(t, GroupCacheStats{Requests: 4, CachedRequests: 3, InputTokens: 600, CachedTokens: 170}, stats["renamed"])
|
|
assert.Equal(t, GroupCacheStats{Requests: 1, InputTokens: 50}, stats["automatic-down"])
|
|
assert.Equal(t, GroupCacheStats{}, stats["unused"])
|
|
assert.NotContains(t, stats, "manual-off")
|
|
})
|
|
}
|
|
}
|