feat: 完善模型请求适配与输出限制
This commit is contained in:
@@ -0,0 +1,194 @@
|
||||
package httpapi
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func TestKeling30TurboSimulationFlow(t *testing.T) {
|
||||
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
|
||||
if databaseURL == "" {
|
||||
t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run the Kling simulation integration flow")
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
applyMigration(t, ctx, databaseURL)
|
||||
db, err := store.Connect(ctx, databaseURL)
|
||||
if err != nil {
|
||||
t.Fatalf("connect store: %v", err)
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
serverCtx, cancelServer := context.WithCancel(ctx)
|
||||
defer cancelServer()
|
||||
server := httptest.NewServer(NewServerWithContext(serverCtx, config.Config{
|
||||
AppEnv: "test",
|
||||
HTTPAddr: ":0",
|
||||
DatabaseURL: databaseURL,
|
||||
IdentityMode: "hybrid",
|
||||
JWTSecret: "test-secret",
|
||||
CORSAllowedOrigin: "*",
|
||||
}, db, slog.New(slog.NewTextHandler(io.Discard, nil))))
|
||||
defer server.Close()
|
||||
|
||||
suffix := strconv.FormatInt(time.Now().UnixNano(), 10)
|
||||
username := "kling_sim_" + suffix
|
||||
password := "password123"
|
||||
var registerResponse struct {
|
||||
AccessToken string `json:"accessToken"`
|
||||
}
|
||||
doJSON(t, server.URL, http.MethodPost, "/api/v1/auth/register", "", map[string]any{
|
||||
"username": username,
|
||||
"email": username + "@example.com",
|
||||
"password": password,
|
||||
}, http.StatusCreated, ®isterResponse)
|
||||
|
||||
var apiKeyResponse struct {
|
||||
Secret string `json:"secret"`
|
||||
}
|
||||
doJSON(t, server.URL, http.MethodPost, "/api/v1/api-keys", registerResponse.AccessToken, map[string]any{
|
||||
"name": "kling simulation key",
|
||||
}, http.StatusCreated, &apiKeyResponse)
|
||||
|
||||
if _, err := db.Pool().Exec(ctx, `
|
||||
UPDATE gateway_users
|
||||
SET roles = '["admin"]'::jsonb
|
||||
WHERE username = $1`, username); err != nil {
|
||||
t.Fatalf("promote Kling simulation user: %v", err)
|
||||
}
|
||||
var loginResponse struct {
|
||||
AccessToken string `json:"accessToken"`
|
||||
}
|
||||
doJSON(t, server.URL, http.MethodPost, "/api/v1/auth/login", "", map[string]any{
|
||||
"account": username,
|
||||
"password": password,
|
||||
}, http.StatusOK, &loginResponse)
|
||||
|
||||
var gatewayUserID string
|
||||
if err := db.Pool().QueryRow(
|
||||
ctx,
|
||||
`SELECT id::text FROM gateway_users WHERE username = $1`,
|
||||
username,
|
||||
).Scan(&gatewayUserID); err != nil {
|
||||
t.Fatalf("read Kling simulation user id: %v", err)
|
||||
}
|
||||
doJSON(t, server.URL, http.MethodPatch, "/api/admin/users/"+gatewayUserID+"/wallet", loginResponse.AccessToken, map[string]any{
|
||||
"currency": "resource",
|
||||
"balance": 1000,
|
||||
"reason": "seed Kling simulation wallet",
|
||||
}, http.StatusOK, nil)
|
||||
|
||||
var platform struct {
|
||||
ID string `json:"id"`
|
||||
}
|
||||
doJSON(t, server.URL, http.MethodPost, "/api/admin/platforms", loginResponse.AccessToken, map[string]any{
|
||||
"provider": "keling",
|
||||
"platformKey": "keling-simulation-" + suffix,
|
||||
"name": "Kling Simulation",
|
||||
"baseUrl": "https://api-beijing.klingai.com/v1",
|
||||
"authType": "AccessKey-SecretKey",
|
||||
"credentials": map[string]any{
|
||||
"accessKey": "legacy-ak",
|
||||
"secretKey": "legacy-sk",
|
||||
},
|
||||
}, http.StatusCreated, &platform)
|
||||
|
||||
doJSON(t, server.URL, http.MethodPost, "/api/admin/platforms/"+platform.ID+"/models", loginResponse.AccessToken, map[string]any{
|
||||
"canonicalModelKey": "keling:kling-3.0-turbo",
|
||||
"modelName": "kling-3.0-turbo",
|
||||
"providerModelName": "kling-3.0-turbo",
|
||||
"modelAlias": "可灵3.0 Turbo",
|
||||
"modelType": []string{"video_generate", "image_to_video"},
|
||||
"displayName": "可灵3.0 Turbo",
|
||||
}, http.StatusCreated, nil)
|
||||
|
||||
assertKeling30TurboSimulationTask := func(
|
||||
name string,
|
||||
request map[string]any,
|
||||
expectedModelType string,
|
||||
) {
|
||||
t.Helper()
|
||||
t.Run(name, func(t *testing.T) {
|
||||
var response struct {
|
||||
Task struct {
|
||||
ID string `json:"id"`
|
||||
Status string `json:"status"`
|
||||
RunMode string `json:"runMode"`
|
||||
ModelType string `json:"modelType"`
|
||||
ResolvedModel string `json:"resolvedModel"`
|
||||
Result map[string]any `json:"result"`
|
||||
Metrics map[string]any `json:"metrics"`
|
||||
BillingSummary map[string]any `json:"billingSummary"`
|
||||
FinalChargeAmount float64 `json:"finalChargeAmount"`
|
||||
ResponseDurationMS int64 `json:"responseDurationMs"`
|
||||
} `json:"task"`
|
||||
}
|
||||
doJSON(
|
||||
t,
|
||||
server.URL,
|
||||
http.MethodPost,
|
||||
"/api/v1/videos/generations",
|
||||
apiKeyResponse.Secret,
|
||||
request,
|
||||
http.StatusAccepted,
|
||||
&response,
|
||||
)
|
||||
|
||||
task := response.Task
|
||||
if task.ID == "" ||
|
||||
task.Status != "succeeded" ||
|
||||
task.RunMode != "simulation" ||
|
||||
task.ModelType != expectedModelType ||
|
||||
task.ResolvedModel != "kling-3.0-turbo" {
|
||||
t.Fatalf("unexpected Kling simulation task: %+v", task)
|
||||
}
|
||||
data, _ := task.Result["data"].([]any)
|
||||
item, _ := data[0].(map[string]any)
|
||||
if item["video_url"] != "/static/simulation/video.mp4" ||
|
||||
item["assetSource"] != "simulation" {
|
||||
t.Fatalf("unexpected Kling simulation result: %+v", task.Result)
|
||||
}
|
||||
if task.FinalChargeAmount <= 0 ||
|
||||
task.BillingSummary["finalCharge"] == nil ||
|
||||
task.Metrics["parameterPreprocessingSummary"] == nil ||
|
||||
task.ResponseDurationMS <= 0 {
|
||||
t.Fatalf("Kling simulation should preserve billing, preprocessing and timing: %+v", task)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
assertKeling30TurboSimulationTask("text-to-video", map[string]any{
|
||||
"model": "可灵3.0 Turbo",
|
||||
"runMode": "simulation",
|
||||
"simulation": true,
|
||||
"simulationDurationMs": 5,
|
||||
"prompt": "A cinematic city reveal",
|
||||
"duration": 8,
|
||||
"resolution": "1080p",
|
||||
"aspect_ratio": "9:16",
|
||||
"audio": true,
|
||||
}, "video_generate")
|
||||
|
||||
assertKeling30TurboSimulationTask("image-to-video", map[string]any{
|
||||
"model": "可灵3.0 Turbo",
|
||||
"runMode": "simulation",
|
||||
"simulation": true,
|
||||
"simulationDurationMs": 5,
|
||||
"prompt": "The subject looks toward the camera",
|
||||
"image": "https://example.com/first.png",
|
||||
"duration": 5,
|
||||
"resolution": "720p",
|
||||
"audio": true,
|
||||
}, "image_to_video")
|
||||
}
|
||||
Reference in New Issue
Block a user