From 69057a1b9574a3b33de4645240cdc6e87140ec81 Mon Sep 17 00:00:00 2001 From: ila Date: Sun, 23 Aug 2026 23:44:23 +0800 Subject: [PATCH] =?UTF-8?q?test:=20=E8=A1=A5=E5=85=85=20MVP-1=20=E5=A4=9A?= =?UTF-8?q?=E5=8D=8F=E8=AE=AE=E9=9B=86=E6=88=90=E9=AA=8C=E6=94=B6=20(#27)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- portal/worker/mysql_integration_test.go | 182 ++++++++++++++++++++++++ 1 file changed, 182 insertions(+) diff --git a/portal/worker/mysql_integration_test.go b/portal/worker/mysql_integration_test.go index cf6a8f5..76bf3df 100644 --- a/portal/worker/mysql_integration_test.go +++ b/portal/worker/mysql_integration_test.go @@ -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 == "" {