67 lines
2.3 KiB
Go
67 lines
2.3 KiB
Go
package workbench
|
|
|
|
import (
|
|
"testing"
|
|
|
|
"yingce/backend/internal/model"
|
|
)
|
|
|
|
func TestValidateTaskCapabilitiesInputContract(t *testing.T) {
|
|
capability := map[string]any{
|
|
"referenceKinds": []any{"image", "video"},
|
|
"inputContract": map[string]any{
|
|
"promptRequired": false,
|
|
"slots": []any{
|
|
map[string]any{
|
|
"key": "subject",
|
|
"name": "Subject",
|
|
"mediaTypes": []any{"image"},
|
|
"minCount": float64(1),
|
|
"maxCount": float64(2),
|
|
"required": true,
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
tests := []struct {
|
|
name string
|
|
refs []TaskReference
|
|
want bool
|
|
}{
|
|
{name: "valid slot reference", refs: []TaskReference{{ResourceID: "image-1", Kind: "image", SlotKey: "subject"}}, want: true},
|
|
{name: "missing required slot", want: false},
|
|
{name: "unknown slot", refs: []TaskReference{{ResourceID: "image-1", Kind: "image", SlotKey: "unknown"}}, want: false},
|
|
{name: "media type mismatch", refs: []TaskReference{{ResourceID: "video-1", Kind: "video", SlotKey: "subject"}}, want: false},
|
|
{name: "slot count exceeded", refs: []TaskReference{{ResourceID: "image-1", Kind: "image", SlotKey: "subject"}, {ResourceID: "image-2", Kind: "image", SlotKey: "subject"}, {ResourceID: "image-3", Kind: "image", SlotKey: "subject"}}, want: false},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
err := validateTaskCapabilities(KindVideo, capability, test.refs)
|
|
if (err == nil) != test.want {
|
|
t.Fatalf("validateTaskCapabilities() error = %v, want success = %t", err, test.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestValidateTaskCapabilitiesLegacyWithoutSlots(t *testing.T) {
|
|
capability := map[string]any{"referenceKinds": []any{"image"}}
|
|
refs := []TaskReference{{ResourceID: "image-1", Kind: "image"}}
|
|
if err := validateTaskCapabilities(KindImage, capability, refs); err != nil {
|
|
t.Fatalf("legacy reference without slot was rejected: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestToolRequiresPromptDefaultsAndOverride(t *testing.T) {
|
|
legacy := &model.WorkbenchTool{CapabilityJSON: `{}`}
|
|
if !toolRequiresPrompt(legacy) {
|
|
t.Fatal("legacy tool must keep requiring a prompt")
|
|
}
|
|
optionalPrompt := &model.WorkbenchTool{CapabilityJSON: `{"inputContract":{"promptRequired":false}}`}
|
|
if toolRequiresPrompt(optionalPrompt) {
|
|
t.Fatal("tool configured for optional prompts must allow an empty prompt")
|
|
}
|
|
}
|