cneicy
a5f7b8e586
Publish Docker image (Multi-arch) / Build & push (amd64) (push) Failing after 1m40s
Release (Linux, macOS, Windows) / Linux Release (push) Failing after 42m16s
Publish Docker image (Multi-arch) / Build & push (arm64) (push) Canceled after 0s
Release (Linux, macOS, Windows) / macOS Release (push) Canceled after 0s
Release (Linux, macOS, Windows) / Windows Release (push) Canceled after 0s
Publish Docker image (Multi-arch) / Create multi-arch manifests (push) Canceled after 0s
Sync production source including routing, billing, Astra detection and public uptime. Update custom build instructions and license packaging.
578 lines
24 KiB
Go
578 lines
24 KiB
Go
package controller
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/constant"
|
|
"github.com/QuantumNous/new-api/dto"
|
|
"github.com/QuantumNous/new-api/model"
|
|
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
|
|
"github.com/QuantumNous/new-api/relay"
|
|
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
|
"github.com/QuantumNous/new-api/service"
|
|
"github.com/QuantumNous/new-api/types"
|
|
"github.com/gin-gonic/gin"
|
|
"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"
|
|
)
|
|
|
|
type taskSubmissionTestBilling struct {
|
|
events *[]string
|
|
settleErr error
|
|
reserveErr error
|
|
onSettle func()
|
|
refunds int
|
|
}
|
|
|
|
func (b *taskSubmissionTestBilling) Settle(int) error {
|
|
*b.events = append(*b.events, "settle")
|
|
if b.onSettle != nil {
|
|
b.onSettle()
|
|
}
|
|
return b.settleErr
|
|
}
|
|
|
|
func (b *taskSubmissionTestBilling) Refund(*gin.Context) {
|
|
*b.events = append(*b.events, "refund")
|
|
b.refunds++
|
|
}
|
|
|
|
func (b *taskSubmissionTestBilling) NeedsRefund() bool { return b.refunds == 0 }
|
|
func (b *taskSubmissionTestBilling) GetPreConsumedQuota() int { return 0 }
|
|
func (b *taskSubmissionTestBilling) Reserve(int) error {
|
|
*b.events = append(*b.events, "reserve")
|
|
return b.reserveErr
|
|
}
|
|
|
|
func TestPresentTaskSubmissionUsesNativePresenterAfterPersistence(t *testing.T) {
|
|
plugin, err := pluginruntime.CompilePlugin(`
|
|
export const meta = {apiVersion:1,key:"presenter-test",name:"Presenter",version:"1.0.0",author:{name:"Test"},models:["model"],fetchMode:"per_task",routes:[{method:"POST",path:"/vendor/jobs",type:"submit",decode:"decode",render:"created"}]};
|
|
export const native = {decode:function(ctx){return {kind:"submit",model:"model",requestBody:ctx.body.value};},created:function(ctx,task){return {data:{task_id:task.task_id},upstream:task.data};}};
|
|
export function buildSubmitRequest(){return {}} export function parseSubmitResponse(){return {taskId:"upstream"}} export function buildQueryRequest(){return {}} export function parseTaskResult(){return {status:"SUCCESS"}}
|
|
`, pluginruntime.Options{})
|
|
require.NoError(t, err)
|
|
priceData := types.PriceData{}
|
|
priceData.AddOtherRatio("seconds", 5)
|
|
task := &model.Task{TaskID: "task_public", SubmitTime: 123}
|
|
task.SetData(map[string]any{"task_id": "upstream_private"})
|
|
outcome := &taskSubmissionOutcome{
|
|
Result: &relay.TaskSubmitResult{},
|
|
Task: task,
|
|
RelayInfo: &relaycommon.RelayInfo{PriceData: priceData},
|
|
}
|
|
recorder := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(recorder)
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/vendor/jobs", strings.NewReader(`{"model":"model"}`))
|
|
c.Set(pluginruntime.ContextKeyPinnedRoute, pluginruntime.PinnedRoute{Plugin: plugin, Route: plugin.Meta.Routes[0]})
|
|
c.Set(pluginruntime.ContextKeyRouteRequest, pluginruntime.RouteRequestContext{Path: "/vendor/jobs", Method: http.MethodPost, Body: map[string]any{"kind": "json", "value": map[string]any{"model": "model"}}})
|
|
|
|
presentTaskSubmission(c, outcome)
|
|
|
|
assert.JSONEq(t, `{
|
|
"data":{"task_id":"task_public"},
|
|
"upstream":{"task_id":"upstream_private"}
|
|
}`, recorder.Body.String())
|
|
assert.JSONEq(t, `{"seconds":5}`, recorder.Header().Get("X-New-Api-Other-Ratios"))
|
|
}
|
|
|
|
func TestPresentTaskSubmissionFallbackUsesPersistedPublicID(t *testing.T) {
|
|
recorder := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(recorder)
|
|
outcome := &taskSubmissionOutcome{
|
|
Result: &relay.TaskSubmitResult{},
|
|
Task: &model.Task{TaskID: "task_persisted", SubmitTime: 456},
|
|
RelayInfo: &relaycommon.RelayInfo{OriginModelName: "video-model"},
|
|
}
|
|
|
|
presentTaskSubmission(c, outcome)
|
|
|
|
assert.JSONEq(t, `{
|
|
"id":"task_persisted",
|
|
"task_id":"task_persisted",
|
|
"status":"queued",
|
|
"model":"video-model",
|
|
"created_at":456
|
|
}`, recorder.Body.String())
|
|
}
|
|
|
|
func TestPresentTaskSubmissionUsesHostOpenAIVideoCreateReceipt(t *testing.T) {
|
|
recorder := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(recorder)
|
|
c.Set(pluginruntime.ContextKeyPinnedEndpoint, pluginruntime.PinnedEndpoint{
|
|
Protocol: "openai_video",
|
|
Operation: pluginruntime.HostProtocolOperation{Name: "create"},
|
|
})
|
|
task := &model.Task{
|
|
TaskID: "task_public",
|
|
Status: model.TaskStatusSubmitted,
|
|
Progress: "0%",
|
|
CreatedAt: 456,
|
|
Properties: model.Properties{OriginModelName: "video-model"},
|
|
}
|
|
outcome := &taskSubmissionOutcome{Result: &relay.TaskSubmitResult{}, Task: task, RelayInfo: &relaycommon.RelayInfo{}}
|
|
|
|
presentTaskSubmission(c, outcome)
|
|
|
|
assert.JSONEq(t, `{"id":"task_public","object":"video","model":"video-model","status":"queued","progress":0,"created_at":456}`, recorder.Body.String())
|
|
assert.NotContains(t, recorder.Body.String(), "task_id")
|
|
}
|
|
|
|
func TestExecuteTaskSubmissionRefundsWhenInsertFails(t *testing.T) {
|
|
events := make([]string, 0, 3)
|
|
database := setupTaskSubmissionDatabase(t, false, &events)
|
|
_ = database
|
|
billing := &taskSubmissionTestBilling{events: &events}
|
|
c := taskSubmissionTestContext()
|
|
info := taskSubmissionRelayInfo(billing)
|
|
|
|
outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) {
|
|
return &relay.TaskSubmitResult{
|
|
UpstreamTaskID: "upstream_private",
|
|
Platform: constant.TaskPlatform("plugin"),
|
|
}, nil
|
|
})
|
|
|
|
assert.Nil(t, outcome)
|
|
require.NotNil(t, taskErr)
|
|
assert.Equal(t, "task_insert_failed", taskErr.Code)
|
|
assert.Equal(t, []string{"reserve", "insert", "refund"}, events)
|
|
assert.Equal(t, 1, billing.refunds)
|
|
assert.False(t, c.Writer.Written())
|
|
}
|
|
|
|
func TestExecuteTaskSubmissionSettlementFailureStaysDurableAndWritesNothing(t *testing.T) {
|
|
events := make([]string, 0, 3)
|
|
database := setupTaskSubmissionDatabase(t, true, &events)
|
|
billing := &taskSubmissionTestBilling{events: &events, settleErr: errors.New("settlement failed")}
|
|
c := taskSubmissionTestContext()
|
|
info := taskSubmissionRelayInfo(billing)
|
|
|
|
outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) {
|
|
return &relay.TaskSubmitResult{
|
|
UpstreamTaskID: "upstream_private",
|
|
Platform: constant.TaskPlatform("plugin"),
|
|
}, nil
|
|
})
|
|
|
|
assert.Nil(t, outcome)
|
|
require.NotNil(t, taskErr)
|
|
assert.Equal(t, "task_billing_settlement_failed", taskErr.Code)
|
|
assert.Equal(t, []string{"reserve", "insert", "settle"}, events)
|
|
assert.Zero(t, billing.refunds)
|
|
var count int64
|
|
require.NoError(t, database.Model(&model.Task{}).Where("task_id = ?", "task_public").Count(&count).Error)
|
|
assert.Equal(t, int64(1), count)
|
|
assert.False(t, c.Writer.Written())
|
|
}
|
|
|
|
func TestExecuteTaskSubmissionPersistsPinnedPluginProvenance(t *testing.T) {
|
|
events := make([]string, 0, 3)
|
|
database := setupTaskSubmissionDatabase(t, true, &events)
|
|
previousLogConsumeEnabled := common.LogConsumeEnabled
|
|
common.LogConsumeEnabled = false
|
|
t.Cleanup(func() { common.LogConsumeEnabled = previousLogConsumeEnabled })
|
|
|
|
c := taskSubmissionTestContext()
|
|
c.Set(common.RequestIdKey, "request-public")
|
|
c.Set(pluginruntime.ContextKeyPinnedPlugin, pluginruntime.PinnedPlugin{
|
|
Generation: &pluginruntime.RoutingGeneration{Number: 42},
|
|
Plugin: &pluginruntime.LoadedPlugin{Meta: pluginruntime.Meta{
|
|
Key: "document-parser",
|
|
Name: "Document Parser",
|
|
Version: "1.2.3",
|
|
APIVersion: 1,
|
|
Author: pluginruntime.AuthorMeta{
|
|
Name: "Community Author",
|
|
URL: "https://plugins.example/author",
|
|
},
|
|
}},
|
|
})
|
|
billing := &taskSubmissionTestBilling{events: &events}
|
|
info := taskSubmissionRelayInfo(billing)
|
|
|
|
outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) {
|
|
return &relay.TaskSubmitResult{
|
|
UpstreamTaskID: "upstream-private",
|
|
Platform: constant.TaskPlatform("document-parser"),
|
|
}, nil
|
|
})
|
|
|
|
require.Nil(t, taskErr)
|
|
require.NotNil(t, outcome)
|
|
require.NotNil(t, outcome.Task.PrivateData.Execution)
|
|
require.NotNil(t, outcome.Task.PrivateData.Execution.TaskPlugin)
|
|
assert.Equal(t, "request-public", outcome.Task.PrivateData.Execution.RequestID)
|
|
assert.Equal(t, "/plugin/submit", outcome.Task.PrivateData.Execution.RequestPath)
|
|
assert.Equal(t, "1.2.3", outcome.Task.PrivateData.Execution.TaskPlugin.Version)
|
|
assert.Equal(t, uint64(42), outcome.Task.PrivateData.Execution.TaskPlugin.Generation)
|
|
require.NotNil(t, outcome.Task.PrivateData.Execution.TaskPlugin.Author)
|
|
assert.Equal(t, "Community Author", outcome.Task.PrivateData.Execution.TaskPlugin.Author.Name)
|
|
assert.Equal(t, "https://plugins.example/author", outcome.Task.PrivateData.Execution.TaskPlugin.Author.URL)
|
|
|
|
var stored model.Task
|
|
require.NoError(t, database.Where("task_id = ?", "task_public").First(&stored).Error)
|
|
require.NotNil(t, stored.PrivateData.Execution)
|
|
require.NotNil(t, stored.PrivateData.Execution.TaskPlugin)
|
|
assert.Equal(t, "document-parser", stored.PrivateData.Execution.TaskPlugin.Key)
|
|
require.NotNil(t, stored.PrivateData.Execution.TaskPlugin.Author)
|
|
assert.Equal(t, "Community Author", stored.PrivateData.Execution.TaskPlugin.Author.Name)
|
|
assert.Equal(t, "upstream-private", stored.PrivateData.UpstreamTaskID)
|
|
}
|
|
|
|
func TestExecuteTaskSubmissionRefundsCancellationBeforeDurableBarrier(t *testing.T) {
|
|
events := make([]string, 0, 2)
|
|
setupTaskSubmissionDatabase(t, true, &events)
|
|
billing := &taskSubmissionTestBilling{events: &events}
|
|
c := taskSubmissionTestContext()
|
|
requestContext, cancel := context.WithCancel(c.Request.Context())
|
|
c.Request = c.Request.WithContext(requestContext)
|
|
info := taskSubmissionRelayInfo(billing)
|
|
|
|
outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) {
|
|
cancel()
|
|
return &relay.TaskSubmitResult{
|
|
UpstreamTaskID: "upstream_private",
|
|
Platform: constant.TaskPlatform("plugin"),
|
|
}, nil
|
|
})
|
|
|
|
assert.Nil(t, outcome)
|
|
require.NotNil(t, taskErr)
|
|
assert.Equal(t, "request_cancelled", taskErr.Code)
|
|
assert.Equal(t, []string{"refund"}, events)
|
|
assert.Equal(t, 1, billing.refunds)
|
|
assert.False(t, c.Writer.Written())
|
|
}
|
|
|
|
func TestExecuteTaskSubmissionDisconnectBeforeUpstreamAcceptanceSkipsSubmitAndRefunds(t *testing.T) {
|
|
events := make([]string, 0, 1)
|
|
setupTaskSubmissionDatabase(t, true, &events)
|
|
billing := &taskSubmissionTestBilling{events: &events}
|
|
c := taskSubmissionTestContext()
|
|
requestContext, cancel := context.WithCancel(c.Request.Context())
|
|
cancel()
|
|
c.Request = c.Request.WithContext(requestContext)
|
|
info := taskSubmissionRelayInfo(billing)
|
|
submitted := false
|
|
|
|
outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) {
|
|
submitted = true
|
|
return nil, nil
|
|
})
|
|
|
|
assert.Nil(t, outcome)
|
|
require.NotNil(t, taskErr)
|
|
assert.Equal(t, "request_cancelled", taskErr.Code)
|
|
assert.False(t, submitted)
|
|
assert.Equal(t, []string{"refund"}, events)
|
|
assert.Equal(t, 1, billing.refunds)
|
|
assert.False(t, c.Writer.Written())
|
|
}
|
|
|
|
func TestExecuteTaskSubmissionCallerCancellationDuringSubmitRefundsBeforeDurableBarrier(t *testing.T) {
|
|
events := make([]string, 0, 1)
|
|
setupTaskSubmissionDatabase(t, true, &events)
|
|
billing := &taskSubmissionTestBilling{events: &events}
|
|
c := taskSubmissionTestContext()
|
|
requestContext, cancel := context.WithCancel(c.Request.Context())
|
|
c.Request = c.Request.WithContext(requestContext)
|
|
info := taskSubmissionRelayInfo(billing)
|
|
submitStarted := make(chan struct{})
|
|
done := make(chan struct{})
|
|
var outcome *taskSubmissionOutcome
|
|
var taskErr *dto.TaskError
|
|
|
|
go func() {
|
|
defer close(done)
|
|
outcome, taskErr = executeTaskSubmissionWith(c, info, func(c *gin.Context, _ *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) {
|
|
close(submitStarted)
|
|
<-c.Request.Context().Done()
|
|
return nil, service.TaskErrorWrapperLocal(c.Request.Context().Err(), "do_request_failed", http.StatusInternalServerError)
|
|
})
|
|
}()
|
|
select {
|
|
case <-submitStarted:
|
|
case <-time.After(2 * time.Second):
|
|
require.FailNow(t, "submission did not start")
|
|
}
|
|
cancel()
|
|
select {
|
|
case <-done:
|
|
case <-time.After(2 * time.Second):
|
|
require.FailNow(t, "submission did not stop after disconnect")
|
|
}
|
|
|
|
assert.Nil(t, outcome)
|
|
require.NotNil(t, taskErr)
|
|
assert.Equal(t, "request_cancelled", taskErr.Code)
|
|
assert.Equal(t, []string{"refund"}, events)
|
|
assert.Equal(t, 1, billing.refunds)
|
|
assert.False(t, c.Writer.Written())
|
|
}
|
|
|
|
func TestExecuteTaskSubmissionDisconnectAfterDurableInsertDoesNotRefund(t *testing.T) {
|
|
events := make([]string, 0, 3)
|
|
database := setupTaskSubmissionDatabase(t, true, &events)
|
|
previousLogConsumeEnabled := common.LogConsumeEnabled
|
|
common.LogConsumeEnabled = false
|
|
t.Cleanup(func() { common.LogConsumeEnabled = previousLogConsumeEnabled })
|
|
c := taskSubmissionTestContext()
|
|
requestContext, cancel := context.WithCancel(c.Request.Context())
|
|
c.Request = c.Request.WithContext(requestContext)
|
|
billing := &taskSubmissionTestBilling{
|
|
events: &events,
|
|
onSettle: cancel,
|
|
}
|
|
info := taskSubmissionRelayInfo(billing)
|
|
|
|
outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) {
|
|
return &relay.TaskSubmitResult{
|
|
UpstreamTaskID: "upstream_private",
|
|
Platform: constant.TaskPlatform("plugin"),
|
|
}, nil
|
|
})
|
|
|
|
require.Nil(t, taskErr)
|
|
require.NotNil(t, outcome)
|
|
assert.Equal(t, "task_public", outcome.Task.TaskID)
|
|
assert.Equal(t, []string{"reserve", "insert", "settle"}, events)
|
|
assert.Zero(t, billing.refunds)
|
|
var count int64
|
|
require.NoError(t, database.Model(&model.Task{}).Where("task_id = ?", "task_public").Count(&count).Error)
|
|
assert.Equal(t, int64(1), count)
|
|
assert.False(t, c.Writer.Written())
|
|
}
|
|
|
|
func setupTaskSubmissionDatabase(t *testing.T, migrate bool, events *[]string) *gorm.DB {
|
|
t.Helper()
|
|
previousDB := model.DB
|
|
database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
|
require.NoError(t, err)
|
|
require.NoError(t, database.Callback().Create().Before("gorm:create").Register("test:task-submit-order", func(*gorm.DB) {
|
|
*events = append(*events, "insert")
|
|
}))
|
|
if migrate {
|
|
require.NoError(t, database.AutoMigrate(&model.Task{}))
|
|
}
|
|
model.DB = database
|
|
t.Cleanup(func() { model.DB = previousDB })
|
|
return database
|
|
}
|
|
|
|
func taskSubmissionTestContext() *gin.Context {
|
|
recorder := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(recorder)
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/plugin/submit", strings.NewReader(`{}`))
|
|
return c
|
|
}
|
|
|
|
func taskSubmissionRelayInfo(billing relaycommon.BillingSettler) *relaycommon.RelayInfo {
|
|
return &relaycommon.RelayInfo{
|
|
UserId: 1,
|
|
UsingGroup: "default",
|
|
OriginModelName: "plugin-model",
|
|
Billing: billing,
|
|
TaskRelayInfo: &relaycommon.TaskRelayInfo{
|
|
PublicTaskID: "task_public",
|
|
LockedChannel: &model.Channel{Id: 1, Type: constant.ChannelTypeTaskPlugin, Name: "plugin"},
|
|
},
|
|
ChannelMeta: &relaycommon.ChannelMeta{ChannelId: 1, ChannelType: constant.ChannelTypeTaskPlugin},
|
|
}
|
|
}
|
|
|
|
// Uses the real submit adaptor, expression evaluator, BillingSession, task row
|
|
// and consume log. Set TEST_TASK_DB_DIALECT plus TEST_MYSQL_DSN or
|
|
// TEST_POSTGRES_DSN to exercise the same contract on an external test database.
|
|
// Unique table prefixes keep the fixture isolated from all existing tables.
|
|
func TestImmediateTaskSettlementDatabase(t *testing.T) {
|
|
dialect := common.DatabaseType(os.Getenv("TEST_TASK_DB_DIALECT"))
|
|
var driver gorm.Dialector
|
|
switch dialect {
|
|
case "", common.DatabaseTypeSQLite:
|
|
dialect = common.DatabaseTypeSQLite
|
|
driver = sqlite.Open(":memory:")
|
|
case common.DatabaseTypeMySQL:
|
|
require.NotEmpty(t, os.Getenv("TEST_MYSQL_DSN"))
|
|
driver = mysql.Open(os.Getenv("TEST_MYSQL_DSN"))
|
|
case common.DatabaseTypePostgreSQL:
|
|
require.NotEmpty(t, os.Getenv("TEST_POSTGRES_DSN"))
|
|
driver = postgres.New(postgres.Config{DSN: os.Getenv("TEST_POSTGRES_DSN"), PreferSimpleProtocol: true})
|
|
default:
|
|
t.Fatalf("unsupported test dialect %q", dialect)
|
|
}
|
|
db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("tsubmit_%d_", time.Now().UnixNano())}})
|
|
require.NoError(t, err)
|
|
sqlDB, err := db.DB()
|
|
require.NoError(t, err)
|
|
sqlDB.SetMaxOpenConns(1)
|
|
t.Cleanup(func() { _ = sqlDB.Close() })
|
|
models := []any{&model.User{}, &model.Channel{}, &model.Task{}, &model.Log{}}
|
|
require.NoError(t, db.AutoMigrate(models...))
|
|
t.Cleanup(func() { require.NoError(t, db.Migrator().DropTable(models...)) })
|
|
var version string
|
|
if dialect == common.DatabaseTypeSQLite {
|
|
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: %s %s", dialect, version)
|
|
oldDB, oldLogDB := model.DB, model.LOG_DB
|
|
oldMain, oldLog := common.MainDatabaseType(), common.LogDatabaseType()
|
|
oldRedis, oldMemory, oldBatch, oldConsume, oldExport := common.RedisEnabled, common.MemoryCacheEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled, common.DataExportEnabled
|
|
model.DB, model.LOG_DB = db, db
|
|
common.SetDatabaseTypes(dialect, dialect)
|
|
common.RedisEnabled, common.MemoryCacheEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled, common.DataExportEnabled = false, false, false, true, false
|
|
t.Cleanup(func() {
|
|
model.DB, model.LOG_DB = oldDB, oldLogDB
|
|
common.SetDatabaseTypes(oldMain, oldLog)
|
|
common.RedisEnabled, common.MemoryCacheEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled, common.DataExportEnabled = oldRedis, oldMemory, oldBatch, oldConsume, oldExport
|
|
})
|
|
|
|
const expression = `u("units") == 7 ? tier("missing", u("missing") * 1.0) : u("units") == 8 ? tier("invalid", -1.0) : u("units") > 3 ? tier("bulk", u("units") * 0.01) : tier("small", u("units") * 0.01)`
|
|
withTieredBillingConfig(t, map[string]string{"document-model": "tiered_expr"}, map[string]string{"document-model": expression})
|
|
const source = `
|
|
export const meta={apiVersion:1,key:"generic-settlement",name:"Generic settlement",version:"1.0.0",author:{name:"Test"},models:["document-model"],fetchMode:"per_task",usageSchema:{units:{type:"number",unit:"count"}}};
|
|
export function buildSubmitRequest(ctx){return {url:ctx.baseUrl+"/compile",body:ctx.requestBody};}
|
|
export function parseSubmitResponse(ctx,resp){return {taskId:"vendor-job",taskData:resp.body,immediate:{status:resp.body.status,reason:"provider rejected job"}};}
|
|
export function extractUsage(){return {units:4};}
|
|
export function extractUsageOnComplete(ctx,result,body){return body.usage;}
|
|
export function parseTaskResult(){throw new Error("completed submissions must not poll");}
|
|
export function buildQueryRequest(){throw new Error("completed submissions must not poll");}
|
|
`
|
|
plugin, err := pluginruntime.CompilePlugin(source, pluginruntime.Options{})
|
|
require.NoError(t, err)
|
|
for index, tc := range []struct {
|
|
name, status string
|
|
actual any
|
|
count float64
|
|
}{
|
|
{"partial", "SUCCESS", 2, 2}, {"zero", "SUCCESS", 0, 0}, {"larger", "SUCCESS", 6, 6},
|
|
{"invalid usage", "SUCCESS", -1, 4}, {"expression failure", "SUCCESS", 7, 4}, {"negative result", "SUCCESS", 8, 4}, {"missing usage", "SUCCESS", nil, 4}, {"failed", "FAILURE", 9, 0},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
body := map[string]any{"status": tc.status}
|
|
if tc.actual != nil {
|
|
body["usage"] = map[string]any{"units": tc.actual}
|
|
}
|
|
encoded, err := common.Marshal(body)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
_, _ = w.Write(encoded)
|
|
}))
|
|
defer server.Close()
|
|
initial := int(20 * common.QuotaPerUnit)
|
|
user := model.User{Username: fmt.Sprintf("task_user_%d", index), AffCode: fmt.Sprintf("task_aff_%d", index), Quota: initial}
|
|
require.NoError(t, db.Create(&user).Error)
|
|
ch := model.Channel{Name: "test provider", Type: constant.ChannelTypeTaskPlugin}
|
|
require.NoError(t, db.Create(&ch).Error)
|
|
c := taskSubmissionTestContext()
|
|
c.Set("group", "default")
|
|
c.Set("username", user.Username)
|
|
c.Set("task_request", map[string]any{"model": "document-model"})
|
|
c.Set(pluginruntime.ContextKeyPinnedPlugin, pluginruntime.PinnedPlugin{Plugin: plugin})
|
|
common.SetContextKey(c, constant.ContextKeyOriginalModel, "document-model")
|
|
common.SetContextKey(c, constant.ContextKeyChannelBaseUrl, server.URL)
|
|
common.SetContextKey(c, constant.ContextKeyChannelId, ch.Id)
|
|
common.SetContextKey(c, constant.ContextKeyChannelType, ch.Type)
|
|
info := taskSubmissionRelayInfo(nil)
|
|
info.UserId = user.Id
|
|
info.OriginModelName = "document-model"
|
|
info.UserGroup = "default"
|
|
info.UsingGroup = "default"
|
|
info.IsPlayground = true
|
|
info.UserSetting.BillingPreference = "wallet_only"
|
|
info.PublicTaskID = model.GenerateTaskID()
|
|
info.LockedChannel = &ch
|
|
outcome, taskErr := executeTaskSubmission(c, info)
|
|
require.Nil(t, taskErr)
|
|
require.NotNil(t, outcome)
|
|
want := common.QuotaRound(tc.count * 0.01 * common.QuotaPerUnit)
|
|
assert.Equal(t, want, outcome.Result.Quota)
|
|
assert.Equal(t, want, info.PriceData.Quota)
|
|
var stored model.Task
|
|
require.NoError(t, db.Where("task_id = ?", info.PublicTaskID).First(&stored).Error)
|
|
assert.Equal(t, want, stored.Quota)
|
|
assert.Equal(t, model.TaskStatus(tc.status), stored.Status)
|
|
assert.Positive(t, stored.FinishTime)
|
|
assert.Equal(t, float64(4), info.TieredBillingSnapshot.EstimatedQuotaBeforeGroup/(0.01*common.QuotaPerUnit))
|
|
if tc.status == "SUCCESS" {
|
|
assert.Equal(t, tc.count, stored.PrivateData.BillingContext.TieredSnapshot.UsageFacts["units"])
|
|
}
|
|
var updated model.User
|
|
require.NoError(t, db.First(&updated, user.Id).Error)
|
|
assert.Equal(t, initial-want, updated.Quota)
|
|
assert.Equal(t, want, updated.UsedQuota)
|
|
var logs []model.Log
|
|
require.NoError(t, db.Where("user_id = ?", user.Id).Find(&logs).Error)
|
|
require.Len(t, logs, 1)
|
|
assert.Equal(t, want, logs[0].Quota)
|
|
var other map[string]any
|
|
require.NoError(t, common.UnmarshalJsonStr(logs[0].Other, &other))
|
|
if tc.status == "SUCCESS" {
|
|
assert.Equal(t, tc.count, other["usage_facts"].(map[string]any)["units"])
|
|
}
|
|
assert.False(t, c.Writer.Written(), "presentation must follow persistence and settlement")
|
|
require.NoError(t, info.Billing.Settle(want))
|
|
info.Billing.Refund(c)
|
|
require.NoError(t, db.First(&updated, user.Id).Error)
|
|
assert.Equal(t, initial-want, updated.Quota, "terminal settlement is idempotent")
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestAcceptedSubmitStreamNeverRetries(t *testing.T) {
|
|
c := taskSubmissionTestContext()
|
|
assert.Equal(t, service.PolicyDecision{Action: "stop", Reason: "task_accepted", Source: "system"}, decideTaskRetry(c, &dto.TaskError{StatusCode: 502, LocalError: true, NoRetry: true}, 3))
|
|
}
|
|
|
|
// Local task rejections carry a message but no cause; the response and the
|
|
// decision record must still be produced.
|
|
func TestRespondTaskSubmissionErrorWithoutCause(t *testing.T) {
|
|
previousErrorLog := constant.ErrorLogEnabled
|
|
constant.ErrorLogEnabled = false
|
|
t.Cleanup(func() { constant.ErrorLogEnabled = previousErrorLog })
|
|
c := taskSubmissionTestContext()
|
|
taskErr := &dto.TaskError{Code: "get_channel_failed", Message: "no channel", StatusCode: http.StatusServiceUnavailable, LocalError: true}
|
|
require.NotPanics(t, func() { respondTaskSubmissionError(c, taskErr) })
|
|
assert.Equal(t, http.StatusServiceUnavailable, c.Writer.Status())
|
|
events := service.RequestPolicy(c).Events()
|
|
require.Len(t, events, 1)
|
|
assert.Equal(t, service.PolicyDecision{Action: "stop", Reason: "request_failed", Source: "system"}, events[0].Decision)
|
|
assert.Equal(t, http.StatusServiceUnavailable, events[0].Status)
|
|
assert.Equal(t, service.PolicyDecision{Action: "stop", Reason: "local_rejection", Source: "system"}, decideTaskRetry(c, &dto.TaskError{StatusCode: http.StatusForbidden, LocalError: true, Message: "billing"}, 2), "local errors stop once the status rules do not force a retry")
|
|
}
|
|
|
|
func TestExecuteTaskSubmissionRefundsWhenFinalReserveFails(t *testing.T) {
|
|
events := []string{}
|
|
setupTaskSubmissionDatabase(t, true, &events)
|
|
billing := &taskSubmissionTestBilling{events: &events, reserveErr: errors.New("insufficient funds")}
|
|
c := taskSubmissionTestContext()
|
|
info := taskSubmissionRelayInfo(billing)
|
|
outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) {
|
|
return &relay.TaskSubmitResult{Platform: "plugin", Quota: 600, Immediate: &relaycommon.TaskInfo{Status: "SUCCESS"}}, nil
|
|
})
|
|
require.Nil(t, outcome)
|
|
require.NotNil(t, taskErr)
|
|
assert.Equal(t, http.StatusForbidden, taskErr.StatusCode)
|
|
assert.Equal(t, []string{"reserve", "refund"}, events)
|
|
assert.Equal(t, 1, billing.refunds)
|
|
assert.False(t, c.Writer.Written())
|
|
}
|