feat: sync deployed v1.0.0-rc.38-custom.67

model/group_status_test.go

@@ -3,6 +3,7 @@ package model
 import (
 	"context"
 	"fmt"
+	"github.com/QuantumNous/new-api/common"
 	"os"
 	"testing"
 	"time"
@@ -16,6 +17,63 @@ import (
 	"gorm.io/gorm/schema"
 )
 
+func TestGroupModelCostsDatabaseMatrix(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":
+				if os.Getenv("TEST_MYSQL_DSN") == "" {
+					t.Skip("no MySQL test database")
+				}
+				driver = mysql.Open(os.Getenv("TEST_MYSQL_DSN"))
+			case "postgres":
+				if os.Getenv("TEST_POSTGRES_DSN") == "" {
+					t.Skip("no PostgreSQL test database")
+				}
+				driver = postgres.Open(os.Getenv("TEST_POSTGRES_DSN"))
+			}
+			db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("cost_%d_", time.Now().UnixNano())}})
+			require.NoError(t, err)
+			oldDB, oldUnit := LOG_DB, common.QuotaPerUnit
+			LOG_DB = db
+			common.QuotaPerUnit = 500000
+			t.Cleanup(func() {
+				LOG_DB = oldDB
+				common.QuotaPerUnit = oldUnit
+				require.NoError(t, db.Migrator().DropTable(&Log{}))
+				sqlDB, e := db.DB()
+				require.NoError(t, e)
+				require.NoError(t, sqlDB.Close())
+			})
+			require.NoError(t, db.AutoMigrate(&Log{}))
+			logs := []Log{
+				{Type: 2, TokenId: 1, CreatedAt: 150, Group: "a", ModelName: "m", Quota: 50, PromptTokens: 100, CompletionTokens: 100, Other: `{"cache_tokens":80}`},
+				{Type: 2, TokenId: 2, CreatedAt: 150, Group: "a", ModelName: "m", Quota: 150, PromptTokens: 200, CompletionTokens: 100, Other: `{"usage_semantic":"anthropic","cache_tokens":100,"cache_creation_tokens":400}`},
+				{Type: 2, TokenId: 3, CreatedAt: 150, Group: "b", ModelName: "m", Quota: 1000, PromptTokens: 200, CompletionTokens: 100, Other: `{"input_tokens_total":900}`},
+				{Type: 2, TokenId: 1, CreatedAt: 150, Group: "a", ModelName: "free", Quota: 0, PromptTokens: 100},
+				{Type: 2, TokenId: 0, CreatedAt: 150, Group: "a", ModelName: "m", Quota: 999999, PromptTokens: 100},
+				{Type: 2, TokenId: 1, CreatedAt: 99, Group: "a", ModelName: "m", Quota: 999999, PromptTokens: 100},
+				{Type: 2, TokenId: 1, CreatedAt: 200, Group: "a", ModelName: "m", Quota: 999999, PromptTokens: 100},
+				{Type: 5, TokenId: 1, CreatedAt: 150, Group: "a", ModelName: "m", Quota: 999999, PromptTokens: 100},
+				{Type: 2, TokenId: 1, CreatedAt: 150, Group: "a", ModelName: "no-tokens", Quota: 500},
+			}
+			require.NoError(t, db.Create(&logs).Error)
+			result, err := GetGroupModelCosts(context.Background(), time.Unix(100, 0), time.Unix(200, 0))
+			require.NoError(t, err)
+			require.NotNil(t, result["a"]["m"].USDPerMillion)
+			assert.InDelta(t, 0.4, *result["a"]["m"].USDPerMillion, 1e-9)
+			assert.Equal(t, int64(2), result["a"]["m"].Requests)
+			assert.Equal(t, float64(1000), result["a"]["m"].Tokens)
+			assert.InDelta(t, 2, *result["b"]["m"].USDPerMillion, 1e-9)
+			assert.Zero(t, *result["a"]["free"].USDPerMillion)
+			assert.NotContains(t, result["a"], "no-tokens")
+		})
+	}
+}
+
 func TestGroupCacheStatsDatabaseMatrix(t *testing.T) {
 	for _, dialect := range []string{"sqlite", "mysql", "postgres"} {
 		t.Run(dialect, func(t *testing.T) {