125 lines
4.3 KiB
Go
125 lines
4.3 KiB
Go
package control
|
|
|
|
import (
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"os"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"gorm.io/gorm"
|
|
|
|
"gpt-load/internal/channel"
|
|
"gpt-load/internal/storage/models"
|
|
)
|
|
|
|
// TestExternalDatabaseReservedIdentifierQueries verifies that runtime query
|
|
// scopes quote table and column names which are reserved by supported drivers.
|
|
func TestExternalDatabaseReservedIdentifierQueries(t *testing.T) {
|
|
// 不标记 t.Parallel()：依赖 GPT_LOAD_DATABASE_TEST_DSN 的共享外部数据库。
|
|
dsn := strings.TrimSpace(os.Getenv("GPT_LOAD_DATABASE_TEST_DSN"))
|
|
if dsn == "" {
|
|
t.Skip("GPT_LOAD_DATABASE_TEST_DSN is not set")
|
|
}
|
|
db := openControlTestDBWithDSN(t, dsn)
|
|
|
|
if err := homeCredentialRowsScope(db).Find(&[]homeCredentialRow{}).Error; err != nil {
|
|
t.Fatalf("query home credentials: %v", err)
|
|
}
|
|
var setting models.SystemSetting
|
|
if err := globalProxyConfigScope(db).Take(&setting).Error; err != nil &&
|
|
!errors.Is(err, gorm.ErrRecordNotFound) {
|
|
t.Fatalf("query global proxy config: %v", err)
|
|
}
|
|
}
|
|
|
|
// TestExternalDatabaseAccessKeyCostLimitPeriodPermutation verifies that the
|
|
// retained-rule two-phase period move obeys the real MySQL/PostgreSQL unique
|
|
// index while preserving IDs and resetting each changed revision.
|
|
func TestExternalDatabaseAccessKeyCostLimitPeriodPermutation(t *testing.T) {
|
|
// 不标记 t.Parallel()：依赖 GPT_LOAD_DATABASE_TEST_DSN 的共享外部数据库，并发执行有唯一索引冲突等正确性风险。
|
|
dsn := strings.TrimSpace(os.Getenv("GPT_LOAD_DATABASE_TEST_DSN"))
|
|
if dsn == "" {
|
|
t.Skip("GPT_LOAD_DATABASE_TEST_DSN is not set")
|
|
}
|
|
assertAccessKeyCostLimitPeriodPermutation(
|
|
t,
|
|
newServiceFixtureWithDSN(t, dsn),
|
|
[]int64{300, 600, 900},
|
|
[]int64{600, 900, 300},
|
|
)
|
|
}
|
|
|
|
// TestExternalDatabaseGroupPriceReconciliation verifies the actual control
|
|
// write chain on each supported external driver. Two Groups can reference one
|
|
// global model price, and removing the final reference cleans only the
|
|
// automatic row.
|
|
func TestExternalDatabaseGroupPriceReconciliation(t *testing.T) {
|
|
// 不标记 t.Parallel()：依赖 GPT_LOAD_DATABASE_TEST_DSN 的共享外部数据库，并发执行有唯一索引冲突等正确性风险。
|
|
dsn := strings.TrimSpace(os.Getenv("GPT_LOAD_DATABASE_TEST_DSN"))
|
|
if dsn == "" {
|
|
t.Skip("GPT_LOAD_DATABASE_TEST_DSN is not set")
|
|
}
|
|
|
|
fixture := newServiceFixtureWithDSN(t, dsn)
|
|
suffix := time.Now().UnixNano()
|
|
modelID := fmt.Sprintf("external-control-model-%d", suffix)
|
|
create := func(name, upstreamURL string) GroupCreateResult {
|
|
result, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{
|
|
Name: &name,
|
|
ChannelID: channel.OpenAICompatible,
|
|
Params: json.RawMessage(`{"base_url":"` + upstreamURL + `"}`),
|
|
Models: optionalGroupModels{Set: true, Values: []GroupModel{{
|
|
ID: modelID,
|
|
}}},
|
|
Credentials: fmt.Sprintf("sk-external-control-%d", suffix), ConnectionType: "api_key",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("CreateGroup(%q) error = %v", name, err)
|
|
}
|
|
return result
|
|
}
|
|
|
|
first := create(
|
|
fmt.Sprintf("external-control-first-%d", suffix),
|
|
fmt.Sprintf("https://external-control-first-%d.example.com/v1", suffix),
|
|
)
|
|
assertExternalAutomaticModelPriceCount(t, fixture, modelID, 1)
|
|
|
|
second := create(
|
|
fmt.Sprintf("external-control-second-%d", suffix),
|
|
fmt.Sprintf("https://external-control-second-%d.example.com/v1", suffix),
|
|
)
|
|
assertExternalAutomaticModelPriceCount(t, fixture, modelID, 1)
|
|
|
|
empty := GroupModelsUpdateRequest{Models: optionalGroupModels{Set: true, Values: []GroupModel{}}}
|
|
if _, err := fixture.service.UpdateGroupModels(t.Context(), first.GroupID, empty); err != nil {
|
|
t.Fatalf("UpdateGroupModels(first) error = %v", err)
|
|
}
|
|
assertExternalAutomaticModelPriceCount(t, fixture, modelID, 1)
|
|
if _, err := fixture.service.UpdateGroupModels(t.Context(), second.GroupID, empty); err != nil {
|
|
t.Fatalf("UpdateGroupModels(second) error = %v", err)
|
|
}
|
|
assertExternalAutomaticModelPriceCount(t, fixture, modelID, 0)
|
|
}
|
|
|
|
func assertExternalAutomaticModelPriceCount(
|
|
t *testing.T,
|
|
fixture serviceFixture,
|
|
modelID string,
|
|
want int64,
|
|
) {
|
|
t.Helper()
|
|
var count int64
|
|
if err := fixture.db.Model(&models.ModelPrice{}).
|
|
Where("model_id = ? AND is_manual = ?", modelID, false).
|
|
Count(&count).Error; err != nil {
|
|
t.Fatalf("count automatic model price: %v", err)
|
|
}
|
|
if count != want {
|
|
t.Fatalf("automatic price count for %q = %d, want %d", modelID, count, want)
|
|
}
|
|
}
|