127 lines
6.8 KiB
Go
127 lines
6.8 KiB
Go
package protocol
|
|
|
|
import (
|
|
"context"
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
func cangyuanAdapter(t *testing.T, provider string) Adapter {
|
|
t.Helper()
|
|
return officialPackageAdapter(t, "cangyuan-midjourney.yingce-plugin", provider)
|
|
}
|
|
|
|
func TestCangyuanMidjourneyV7RequestAndEditPolling(t *testing.T) {
|
|
adapter := cangyuanAdapter(t, "cangyuan-midjourney-v7")
|
|
request := GenerationRequest{Model: "midjourney-v7", Prompt: "生成一张氛围感美女照片", ImageCount: 1, AspectRatio: "auto", Quality: "high", Resolution: "2k", Extra: map[string]any{"transparentBackground": "true"}}
|
|
spec, err := adapter.BuildCreate(context.Background(), RequestContext{Request: request})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
body := manifestTestBody(t, spec)
|
|
if spec.Path != "/v1/images/generations" || body["n"] != float64(1) || body["response_format"] != "url" || body["async"] != false {
|
|
t.Fatalf("create = %#v, body = %#v", spec, body)
|
|
}
|
|
for _, key := range []string{"size", "images", "quality", "resolution", "background", "speed", "output_format", "mask"} {
|
|
if _, ok := body[key]; ok {
|
|
t.Fatalf("unsupported field %s was emitted: %#v", key, body)
|
|
}
|
|
}
|
|
request.Images = []MediaReference{{URL: "https://example.com/second.png", Role: "edit_source", Order: 2}, {URL: "https://example.com/first.png", Role: "edit_source", Order: 1}}
|
|
request.AspectRatio = "9:16"
|
|
spec, err = adapter.BuildCreate(context.Background(), RequestContext{Request: request})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
body = manifestTestBody(t, spec)
|
|
images := body["images"].([]any)
|
|
if spec.Path != "/v1/images/edits" || body["size"] != "9:16" || len(images) != 2 || images[0] != "https://example.com/first.png" {
|
|
t.Fatalf("edit = %#v, body = %#v", spec, body)
|
|
}
|
|
poll, err := adapter.BuildPoll(context.Background(), PollContext{Request: request, TaskID: "task-42"})
|
|
if err != nil || poll.Path != "/v1/images/edits/task-42" {
|
|
t.Fatalf("edit poll = %#v, err = %v", poll, err)
|
|
}
|
|
}
|
|
|
|
func TestCangyuanMidjourneyV82VersionsReferencesAndMask(t *testing.T) {
|
|
adapter := cangyuanAdapter(t, "cangyuan-midjourney-v82")
|
|
for _, model := range []string{"midjourney-1k", "midjourney-2k"} {
|
|
request := GenerationRequest{Model: model, Prompt: "portrait --v 8.2", ImageCount: 1, Quality: "high"}
|
|
spec, err := adapter.BuildCreate(context.Background(), RequestContext{Request: request})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
body := manifestTestBody(t, spec)
|
|
if body["prompt"] != request.Prompt || body["model"] != model {
|
|
t.Fatalf("version request = %#v", body)
|
|
}
|
|
for _, key := range []string{"quality", "resolution", "background", "output_format", "speed", "reference", "mask"} {
|
|
if _, ok := body[key]; ok {
|
|
t.Fatalf("unexpected %s in %#v", key, body)
|
|
}
|
|
}
|
|
}
|
|
request := GenerationRequest{Model: "midjourney-1k", Prompt: "portrait --v 8.2", Images: []MediaReference{{URL: "https://example.com/source.png", Role: "edit_source"}, {URL: "https://example.com/mask.png", Role: "mask"}}}
|
|
spec, err := adapter.BuildCreate(context.Background(), RequestContext{Request: request})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
body := manifestTestBody(t, spec)
|
|
if spec.Path != "/v1/images/edits" || body["reference"] != "editor" || body["mask"] != "https://example.com/mask.png" || len(body["images"].([]any)) != 1 {
|
|
t.Fatalf("masked edit = %#v, body = %#v", spec, body)
|
|
}
|
|
request.Images = request.Images[:1]
|
|
request.ProviderOptions = map[string]map[string]any{"cangyuan-midjourney-v82": {"reference": "style", "speed": "fast"}}
|
|
spec, err = adapter.BuildCreate(context.Background(), RequestContext{Request: request})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
body = manifestTestBody(t, spec)
|
|
if spec.Path != "/v1/images/generations" || body["reference"] != "style" || body["speed"] != "fast" {
|
|
t.Fatalf("style = %#v, body = %#v", spec, body)
|
|
}
|
|
}
|
|
|
|
func TestCangyuanMidjourneyRejectsInvalidInputs(t *testing.T) {
|
|
for _, tc := range []struct {
|
|
name, provider string
|
|
request GenerationRequest
|
|
}{
|
|
{"model mismatch", "cangyuan-midjourney-v7", GenerationRequest{Model: "midjourney-1k", Prompt: "portrait"}},
|
|
{"invented v82 model", "cangyuan-midjourney-v82", GenerationRequest{Model: "midjourney-v8.2", Prompt: "portrait"}},
|
|
{"multiple submissions", "cangyuan-midjourney-v7", GenerationRequest{Model: "midjourney-v7", Prompt: "portrait", ImageCount: 2}},
|
|
{"long unicode prompt", "cangyuan-midjourney-v7", GenerationRequest{Model: "midjourney-v7", Prompt: strings.Repeat("画", 4001)}},
|
|
{"data URI", "cangyuan-midjourney-v7", GenerationRequest{Model: "midjourney-v7", Prompt: "portrait", Images: []MediaReference{{DataURL: "data:image/png;base64,aA=="}}}},
|
|
{"v7 mask", "cangyuan-midjourney-v7", GenerationRequest{Model: "midjourney-v7", Prompt: "portrait", Images: []MediaReference{{URL: "https://example.com/mask.png", Role: "mask"}}}},
|
|
{"editor missing source", "cangyuan-midjourney-v82", GenerationRequest{Model: "midjourney-1k", Prompt: "portrait", Images: []MediaReference{{URL: "https://example.com/mask.png", Role: "mask"}}}},
|
|
{"invalid speed", "cangyuan-midjourney-v82", GenerationRequest{Model: "midjourney-1k", Prompt: "portrait", ProviderOptions: map[string]map[string]any{"cangyuan-midjourney-v82": {"speed": "turbo"}}}},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
if _, err := cangyuanAdapter(t, tc.provider).BuildCreate(context.Background(), RequestContext{Request: tc.request}); err == nil {
|
|
t.Fatal("invalid input was accepted")
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestCangyuanMidjourneyPreservesAllResultsAndTaskStates(t *testing.T) {
|
|
for _, provider := range []string{"cangyuan-midjourney-v7", "cangyuan-midjourney-v82"} {
|
|
adapter := cangyuanAdapter(t, provider)
|
|
for _, payload := range []string{`{"data":[{"url":"https://example.com/1.png"},{"url":"https://example.com/2.png"},{"url":"https://example.com/3.png"},{"url":"https://example.com/4.png"}]}`, `{"id":"task-42","status":"completed","data":[{"url":"https://example.com/1.png"},{"url":"https://example.com/2.png"},{"url":"https://example.com/3.png"},{"url":"https://example.com/4.png"}]}`} {
|
|
result, err := adapter.ParseCreate(context.Background(), []byte(payload))
|
|
if err != nil || result.Status != StatusSucceeded || result.Result == nil || len(result.Result.Images) != 4 {
|
|
t.Fatalf("%s result = %#v, err = %v", provider, result, err)
|
|
}
|
|
}
|
|
pending, err := adapter.ParseCreate(context.Background(), []byte(`{"id":"task-42","status":"queued"}`))
|
|
if err != nil || pending.TaskID != "task-42" || pending.Status != StatusPending {
|
|
t.Fatalf("pending = %#v, err = %v", pending, err)
|
|
}
|
|
failed, err := adapter.ParsePoll(context.Background(), PollContext{TaskID: "task-42"}, []byte(`{"status":"failed","error":{"code":"invalid_prompt","message":"prompt rejected"}}`))
|
|
if err != nil || failed.TaskID != "task-42" || failed.Status != StatusFailed || failed.Message != "prompt rejected" {
|
|
t.Fatalf("failure = %#v, err = %v", failed, err)
|
|
}
|
|
}
|
|
}
|