test: 补充 MVP-1 多协议集成验收 (#27)
This commit is contained in:
@@ -6,6 +6,9 @@ import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"image"
|
||||
"image/color"
|
||||
"image/png"
|
||||
"net"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
@@ -48,6 +51,9 @@ func TestMySQLWorkerUsesSnapshotAndMockUpstream(t *testing.T) {
|
||||
t.Fatal(tx.Error)
|
||||
}
|
||||
defer tx.Rollback()
|
||||
if err := tx.Model(&model.Generation{}).Where("status IN ?", []model.GenerationStatus{model.StatusPending, model.StatusRunning}).Update("status", model.StatusSucceeded).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
server := httptest.NewServer(mockprovider.Handler{})
|
||||
defer server.Close()
|
||||
@@ -162,6 +168,182 @@ func TestMySQLWorkerUsesSnapshotAndMockUpstream(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestMySQLWorkerImageProtocolsUseMockUpstream(t *testing.T) {
|
||||
testCases := []struct {
|
||||
name string
|
||||
apiType model.APIType
|
||||
capability model.Capability
|
||||
withInput bool
|
||||
}{
|
||||
{name: "images", apiType: model.APIImages, capability: model.CapabilityImageGenerate},
|
||||
{name: "images_edits", apiType: model.APIImagesEdits, capability: model.CapabilityImageEdit, withInput: true},
|
||||
{name: "gemini", apiType: model.APIGemini, capability: model.CapabilityImageEdit, withInput: true},
|
||||
}
|
||||
for _, testCase := range testCases {
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
runMySQLImageProtocolWorker(t, testCase.apiType, testCase.capability, testCase.withInput)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func runMySQLImageProtocolWorker(t *testing.T, apiType model.APIType, capability model.Capability, withInput bool) {
|
||||
t.Helper()
|
||||
dsn := os.Getenv("CHORUS_TEST_DSN")
|
||||
if dsn == "" {
|
||||
t.Skip("CHORUS_TEST_DSN is not set")
|
||||
}
|
||||
db, err := gorm.Open(mysql.Open(dsn), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
sqlDB, err := db.DB()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer sqlDB.Close()
|
||||
tx := db.Begin()
|
||||
if tx.Error != nil {
|
||||
t.Fatal(tx.Error)
|
||||
}
|
||||
defer tx.Rollback()
|
||||
if err := tx.Model(&model.Generation{}).Where("status IN ?", []model.GenerationStatus{model.StatusPending, model.StatusRunning}).Update("status", model.StatusSucceeded).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
server := httptest.NewServer(mockprovider.Handler{})
|
||||
defer server.Close()
|
||||
httpClient, err := safehttp.New(safehttp.Config{
|
||||
Resolver: fixedResolver{}, Timeout: 2 * time.Second, MaxRedirects: 1, AllowedPorts: []uint16{80},
|
||||
DialContext: func(ctx context.Context, network, _ string) (net.Conn, error) {
|
||||
return (&net.Dialer{}).DialContext(ctx, network, server.Listener.Addr().String())
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
factory, err := NewOpenAIFactory(safehttp.NewProviderClient(httpClient), 2<<20)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
catalog, err := NewGORMCatalog(tx)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
runtime, err := NewGORMRuntime(tx)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
queueRepository, err := queue.NewMySQLRepository(tx)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
local, err := platformstorage.NewLocal(platformstorage.Config{Root: t.TempDir(), MaxObjectBytes: 2 << 20, MaxImagePixels: 10000, ThumbnailMaxSide: 32, AllowedImageMIME: map[string]bool{"image/png": true}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
store, err := NewLocalStorage(local)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
suffix := time.Now().UnixNano()
|
||||
user := model.User{Email: fmt.Sprintf("worker-image-%d@chorus.invalid", suffix), PasswordHash: "synthetic", DisplayName: "Worker Image", Status: "active"}
|
||||
providerRow := model.Provider{Slug: fmt.Sprintf("worker-image-%d", suffix), Name: "Worker Image Mock", BaseURL: "http://provider.test/v1", AuthType: "none", Enabled: true}
|
||||
if err := tx.Create(&user).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := tx.Create(&providerRow).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
modelRow := model.ProviderModel{ProviderID: providerRow.ID, Name: "Worker Image", ModelID: "mock-image", APIType: apiType, Kind: model.KindImage, ExtraBody: json.RawMessage("{}"), TimeoutMS: 1000, Weight: 1, Enabled: true}
|
||||
if err := tx.Create(&modelRow).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := tx.Exec("INSERT INTO provider_model_capabilities (provider_model_id, capability) VALUES (?, ?)", modelRow.ID, capability).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
template := model.PromptTemplate{TemplateKey: fmt.Sprintf("worker-image-%d", suffix), Kind: model.KindImage, APIType: apiType, Capability: capability, Name: "Worker Image", Version: 1, TemplateText: "{{.UserPrompt}}", DefaultRoleRule: "", Enabled: true}
|
||||
if err := tx.Create(&template).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
poolSlug := fmt.Sprintf("worker-image-%d", suffix)
|
||||
if err := tx.Exec("INSERT INTO route_pools (slug, name, capability, prompt_template_id, max_failover, version, enabled) VALUES (?, 'Worker Image', ?, ?, 0, 1, TRUE)", poolSlug, capability, template.ID).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var poolID uint64
|
||||
if err := tx.Raw("SELECT id FROM route_pools WHERE slug = ?", poolSlug).Scan(&poolID).Error; err != nil || poolID == 0 {
|
||||
t.Fatalf("pool id=%d error=%v", poolID, err)
|
||||
}
|
||||
if err := tx.Exec("INSERT INTO route_pool_members (route_pool_id, provider_model_id, weight, failure_threshold, open_seconds, half_open_max, enabled) VALUES (?, ?, 1, 2, 60, 1, TRUE)", poolID, modelRow.ID).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var memberID uint64
|
||||
if err := tx.Raw("SELECT id FROM route_pool_members WHERE route_pool_id = ?", poolID).Scan(&memberID).Error; err != nil || memberID == 0 {
|
||||
t.Fatalf("member id=%d error=%v", memberID, err)
|
||||
}
|
||||
|
||||
snapshot := router.RouteSnapshot{Capability: capability, RoutePoolID: poolID, RoutePoolVersion: 1, PromptTemplateID: template.ID, PromptTemplateKey: template.TemplateKey, PromptTemplateVersion: 1, Members: []router.MemberSnapshot{{RoutePoolMemberID: memberID, ProviderModelID: modelRow.ID, Weight: 1, FailureThreshold: 2, OpenSeconds: 60, HalfOpenMax: 1}}}
|
||||
generation := model.Generation{UserID: user.ID, Kind: model.KindImage, IdempotencyKey: fmt.Sprintf("worker-image-%d", suffix), UserPrompt: "integration image", RenderedPrompt: "integration image", CreatedAt: time.Unix(1, 0)}
|
||||
if err := router.ApplySnapshot(&generation, snapshot); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
created, err := queueRepository.CreateIdempotentPrepared(context.Background(), &generation, func(generationID uint64) ([]model.GenerationInput, error) {
|
||||
if !withInput {
|
||||
return nil, nil
|
||||
}
|
||||
objects, putErr := store.PutImage(context.Background(), ImageRequest{Key: fmt.Sprintf("inputs/%d/original", generationID), ThumbnailKey: fmt.Sprintf("inputs/%d/thumbnail", generationID), OwnerID: user.ID, GenerationID: generationID, ContentType: "image/png", Source: bytes.NewReader(workerTestPNG(t))})
|
||||
if putErr != nil {
|
||||
return nil, putErr
|
||||
}
|
||||
return []model.GenerationInput{{Position: 0, Role: model.RolePrimary, OriginalName: "synthetic.png", MIMEType: "image/png", StorageKey: objects.Original.Key, SizeBytes: uint64(objects.Original.Size)}}, nil
|
||||
})
|
||||
if err != nil || !created {
|
||||
t.Fatalf("create=%v error=%v", created, err)
|
||||
}
|
||||
controller := queue.NewController(queueRepository)
|
||||
w, err := New(Config{Owner: "image-integration", LeaseDuration: 5 * time.Second, PollInterval: time.Millisecond, Random: firstRandom{}}, queueRepository, controller, catalog, factory, runtime, store)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
worked, err := w.ProcessOne(context.Background())
|
||||
if err != nil || !worked {
|
||||
t.Fatalf("worked=%v error=%v", worked, err)
|
||||
}
|
||||
var stored model.Generation
|
||||
if err := tx.First(&stored, generation.ID).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if stored.Status != model.StatusSucceeded || stored.ProviderAttemptCount != 1 || !bytes.Contains(stored.Attempts, []byte(`"route_member_id"`)) || !bytes.Contains(stored.Attempts, []byte(`"finished_at"`)) {
|
||||
t.Fatalf("stored generation=%#v attempts=%s", stored, stored.Attempts)
|
||||
}
|
||||
var output model.GenerationOutput
|
||||
if err := tx.Where("generation_id = ?", generation.ID).First(&output).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if output.Kind != model.KindImage || output.StorageKey == nil || output.ThumbnailStorageKey == nil || output.MIMEType == nil || *output.MIMEType != "image/png" {
|
||||
t.Fatalf("stored output=%#v", output)
|
||||
}
|
||||
reader, object, err := store.Open(context.Background(), *output.StorageKey)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer reader.Close()
|
||||
if object.OwnerID != user.ID || object.GenerationID != generation.ID {
|
||||
t.Fatalf("stored object=%#v", object)
|
||||
}
|
||||
}
|
||||
|
||||
func workerTestPNG(t *testing.T) []byte {
|
||||
t.Helper()
|
||||
var encoded bytes.Buffer
|
||||
img := image.NewRGBA(image.Rect(0, 0, 2, 2))
|
||||
img.Set(0, 0, color.RGBA{R: 40, G: 120, B: 200, A: 255})
|
||||
if err := png.Encode(&encoded, img); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return encoded.Bytes()
|
||||
}
|
||||
|
||||
func TestMySQLRuntimeOpensAndLimitsHalfOpenProbe(t *testing.T) {
|
||||
dsn := os.Getenv("CHORUS_TEST_DSN")
|
||||
if dsn == "" {
|
||||
|
||||
Reference in New Issue
Block a user