59 lines
2.1 KiB
Go
59 lines
2.1 KiB
Go
|
package workflowbridge
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"data-terminal-cloud/internal/dify"
|
||
|
|
"data-terminal-cloud/internal/extraction"
|
||
|
|
"net/http/httptest"
|
||
|
|
"testing"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestModelSelectionValidation(t *testing.T) {
|
||
|
|
for _, tc := range []struct {
|
||
|
|
step, model, custom, want string
|
||
|
|
invalid bool
|
||
|
|
}{
|
||
|
|
{"classify", "", "", DefaultTextModel, false},
|
||
|
|
{"write-book", "doubao-seed-2-1-pro-260915", "", "doubao-seed-2-1-pro-260915", false},
|
||
|
|
{"illustrate", "__custom__", "ep-test123", "ep-test123", false},
|
||
|
|
{"daily", "__custom__", "", "", true},
|
||
|
|
{"daily", "https://example.org", "", "", true},
|
||
|
|
{"write-book", DefaultImageModel, "", "", true},
|
||
|
|
{"illustrate", DefaultTextModel, "", "", true},
|
||
|
|
} {
|
||
|
|
got, err := selectedModel(tc.step, tc.model, tc.custom)
|
||
|
|
if got != tc.want || (err != nil) != tc.invalid {
|
||
|
|
t.Errorf("%+v: %q %v", tc, got, err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestIndependentModelsAndRetryRetainsOriginalSelection(t *testing.T) {
|
||
|
|
e := &fixtureEditor{}
|
||
|
|
s := New("secret", e, e, overdueImages{})
|
||
|
|
selected := map[string]string{}
|
||
|
|
s.EditorForModel = func(step, model string) (Editor, error) { selected[step] = model; return e, nil }
|
||
|
|
srv := httptest.NewServer(s)
|
||
|
|
defer srv.Close()
|
||
|
|
b := dify.NewBridge(srv.URL, "secret")
|
||
|
|
ticket, err := b.Register(context.Background(), "extract", extraction.GenerationInput{Content: "[turn-1] 长者：1970年进厂。"})
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
defer b.Forget(ticket)
|
||
|
|
c := s.find(ticket.ID)
|
||
|
|
out := s.step(c, "classify", stepInput{Model: "doubao-seed-2-1-lite-260915"})
|
||
|
|
if out["success"] != true {
|
||
|
|
t.Fatal(out)
|
||
|
|
}
|
||
|
|
s.step(c, "daily", stepInput{Model: DefaultTextModel})
|
||
|
|
retry := s.step(c, "classify", stepInput{Model: DefaultBookModel})
|
||
|
|
if retry["model"] != out["model"] || e.calls.Load() != 2 || selected["classify"] != "doubao-seed-2-1-lite-260915" || selected["daily"] != DefaultTextModel {
|
||
|
|
t.Fatalf("wrong selection: %v %v", selected, retry)
|
||
|
|
}
|
||
|
|
merged := s.step(c, "dedup", stepInput{})
|
||
|
|
if merged["success"] != true || c.Result.ModelRef != "dify:doubao-seed-2-1-lite-260915" {
|
||
|
|
t.Fatalf("actual model attribution lost: %v %q", merged, c.Result.ModelRef)
|
||
|
|
}
|
||
|
|
}
|