cneicy
ceddfd3d6b
Sync production source including routing, billing, Astra detection and public uptime. Update custom build instructions and license packaging.
51 lines
2.3 KiB
Go
51 lines
2.3 KiB
Go
package controller
|
|
|
|
import (
|
|
"context"
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/model"
|
|
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
|
"github.com/QuantumNous/new-api/service"
|
|
"github.com/glebarez/sqlite"
|
|
"github.com/stretchr/testify/require"
|
|
"gorm.io/gorm"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestLocalPelicanGeneratesOnceWithUnknownVerdict(t *testing.T) {
|
|
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
|
require.NoError(t, err)
|
|
conn, err := db.DB()
|
|
require.NoError(t, err)
|
|
conn.SetMaxOpenConns(1)
|
|
old := model.DB
|
|
model.DB = db
|
|
t.Cleanup(func() { model.DB = old; _ = conn.Close() })
|
|
require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.DetectionRun{}, &model.DetectionCall{}, &model.DetectionSchedule{}))
|
|
calls := 0
|
|
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
calls++
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = w.Write([]byte(`{"status":"completed","usage":{"input_tokens":10,"output_tokens":5},"output":[{"type":"message","content":[{"type":"output_text","text":"<html><body>pelican</body></html>"}]}]}`))
|
|
}))
|
|
defer upstream.Close()
|
|
url := upstream.URL
|
|
channel := model.Channel{Type: 1, Key: "test-only", BaseURL: &url, Status: common.ChannelStatusEnabled, Models: "gpt-6-astra"}
|
|
require.NoError(t, db.Create(&channel).Error)
|
|
price, _ := common.Marshal(service.DetectionPrice{Snapshot: billingexpr.BillingSnapshot{ExprString: "p+c", ExprHash: billingexpr.ExprHashString("p+c"), GroupRatio: 1, QuotaPerUnit: 500000, ExprVersion: 1}})
|
|
run := model.DetectionRun{ID: "local-test", RequestKey: "local-test", Kind: "pelican", ChannelID: channel.Id, Status: "running", Deadline: time.Now().Unix() + 60, Reserved: 100, Price: string(price)}
|
|
require.NoError(t, db.Create(&run).Error)
|
|
require.NoError(t, db.Create(&model.DetectionCall{ID: run.ID, RunID: run.ID, SecretHash: model.DetectionSecretHash(pelicanCredential(run.ID)), State: "ready", ExpiresAt: run.Deadline}).Error)
|
|
processLocalPelican(context.Background(), run)
|
|
require.NoError(t, db.First(&run, "id = ?", run.ID).Error)
|
|
require.Equal(t, "succeeded", run.Status)
|
|
require.Equal(t, "unknown", run.Quality)
|
|
require.Contains(t, run.HTML, "pelican")
|
|
require.True(t, run.Settled)
|
|
processLocalPelican(context.Background(), run)
|
|
require.Equal(t, 1, calls)
|
|
}
|