feat: publish v1.0.0-rc.36-custom.1

model/group_status_test.go

@@ -0,0 +1,92 @@
+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")
+		})
+	}
+}