new-api-custom/model/detection_test.go

package model

import (
	"fmt"
	"github.com/QuantumNous/new-api/common"
	"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"
	"os"
	"sync"
	"testing"
	"time"
)

func TestDetectionDatabaseMatrix(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 unset")
				}
				driver = mysql.Open(dsn)
			case "postgres":
				dsn := os.Getenv("TEST_POSTGRES_DSN")
				if dsn == "" {
					t.Skip("TEST_POSTGRES_DSN unset")
				}
				driver = postgres.Open(dsn)
			}
			db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("det_%d_", time.Now().UnixNano())}})
			require.NoError(t, err)
			sqlDB, err := db.DB()
			require.NoError(t, err)
			sqlDB.SetMaxOpenConns(1)
			if dialect != "sqlite" {
				sqlDB.SetMaxOpenConns(4)
			}
			var version string
			if dialect == "sqlite" {
				require.NoError(t, db.Raw("select sqlite_version()").Scan(&version).Error)
			} else {
				require.NoError(t, db.Raw("select version()").Scan(&version).Error)
			}
			t.Logf("database version: %s", version)
			oldDB, oldLogDB, oldRedis := DB, LOG_DB, common.RedisEnabled
			DB = db
			LOG_DB = db
			common.RedisEnabled = false
			t.Cleanup(func() {
				for _, m := range []any{&DetectionCall{}, &DetectionRun{}, &User{}, &Log{}} {
					require.NoError(t, db.Migrator().DropTable(m))
				}
				DB = oldDB
				LOG_DB = oldLogDB
				common.RedisEnabled = oldRedis
				require.NoError(t, sqlDB.Close())
			})
			require.NoError(t, db.AutoMigrate(&Log{}))
			require.NoError(t, db.AutoMigrate(&User{}))
			user := User{Username: "detection-owner", Quota: 1000, Status: common.UserStatusEnabled}
			require.NoError(t, db.Create(&user).Error)
			for range 2 {
				require.NoError(t, db.AutoMigrate(&User{}, &DetectionRun{}, &DetectionCall{}))
			}
			now := time.Now().Unix()
			run := DetectionRun{ID: "detection", RequestKey: "user-idempotency", UserID: user.Id, Kind: "pelican", GroupName: "a", Reserved: 300, CreatedAt: now, Deadline: now + 600, Status: "pending"}
			require.NoError(t, CreateDetection(&run))
			require.NoError(t, CreateDetection(&run))
			require.NoError(t, db.First(&user, user.Id).Error)
			assert.Equal(t, 700, user.Quota)
			second := run
			second.ID = "second"
			second.RequestKey = "second"
			require.Error(t, CreateDetection(&second))
			call := DetectionCall{ID: "call", RunID: run.ID, SecretHash: DetectionSecretHash("secret"), ExpiresAt: now + 900, State: "ready"}
			require.NoError(t, db.Create(&call).Error)
			var wg sync.WaitGroup
			claimedCalls := make(chan *DetectionCall, 2)
			for range 2 {
				wg.Go(func() {
					result, claimErr := ClaimDetectionCall(call.SecretHash)
					if claimErr == nil {
						claimedCalls <- result
					}
				})
			}
			wg.Wait()
			close(claimedCalls)
			require.Len(t, claimedCalls, 1)
			claimed := <-claimedCalls
			_, err = ClaimDetectionCall(call.SecretHash)
			require.Error(t, err)
			require.Error(t, SettleDetection(run.ID))
			claimed.Quota = 42
			claimed.Usage = `{"input_tokens":10,"output_tokens":5}`
			require.NoError(t, FinishDetectionCall(claimed))
			claimed.Quota = 99
			require.NoError(t, FinishDetectionCall(claimed))
			require.NoError(t, SettleDetection(run.ID))
			require.NoError(t, SettleDetection(run.ID))
			require.NoError(t, DeliverDetectionLogs())
			require.NoError(t, DeliverDetectionLogs())
			var logCount int64
			require.NoError(t, db.Model(&Log{}).Count(&logCount).Error)
			assert.EqualValues(t, 1, logCount)
			require.NoError(t, db.First(&user, user.Id).Error)
			assert.Equal(t, 958, user.Quota)
			assert.Equal(t, 42, user.UsedQuota)
			_, err = ClaimDetectionCall(call.SecretHash)
			require.Error(t, err)
			require.NoError(t, db.AutoMigrate(&DetectionRun{}, &DetectionCall{}))
			require.NoError(t, db.AutoMigrate(&DetectionRun{}, &DetectionCall{}))
			var stored DetectionRun
			require.NoError(t, db.First(&stored, "id = ?", run.ID).Error)
			assert.Equal(t, 42, stored.Charged)
			expired := DetectionRun{ID: "expired", RequestKey: "expired", Kind: "pelican", Deadline: now - 1, CreatedAt: now - 901, Status: "pending"}
			require.NoError(t, CreateDetection(&expired))
			expCall := DetectionCall{ID: "expired-call", RunID: expired.ID, SecretHash: DetectionSecretHash("expired"), State: "ready", ExpiresAt: now - 1}
			require.NoError(t, db.Create(&expCall).Error)
			_, err = ClaimDetectionCall(expCall.SecretHash)
			require.Error(t, err)
			require.NoError(t, SettleDetection(expired.ID))
			deleted := User{Username: "deleted-owner", AffCode: "deleted-owner", Quota: 500, Status: common.UserStatusEnabled}
			require.NoError(t, db.Create(&deleted).Error)
			deletedRun := DetectionRun{ID: "deleted", RequestKey: "deleted", UserID: deleted.Id, Kind: "pelican", Reserved: 100, CreatedAt: now}
			require.NoError(t, CreateDetection(&deletedRun))
			require.NoError(t, db.Delete(&deleted).Error)
			require.NoError(t, SettleDetection(deletedRun.ID))
			require.NoError(t, db.Unscoped().First(&deleted, deleted.Id).Error)
			assert.Equal(t, 500, deleted.Quota)
		})
	}
}