96 lines
2.6 KiB
Go
96 lines
2.6 KiB
Go
package vision
|
|
|
|
import (
|
|
"context"
|
|
"strings"
|
|
"testing"
|
|
|
|
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
|
|
)
|
|
|
|
func TestDetectMultimodal(t *testing.T) {
|
|
t.Parallel()
|
|
tests := []struct {
|
|
name string
|
|
msgs []model.Message
|
|
want bool
|
|
}{
|
|
{name: "text only", msgs: []model.Message{model.TextMessage(model.RoleUser, "hi")}, want: false},
|
|
{name: "empty", msgs: nil, want: false},
|
|
{name: "single image", msgs: []model.Message{model.ImageMessage(model.RoleUser, "what", "data:image/png;base64,abc")}, want: true},
|
|
{name: "image no text", msgs: []model.Message{model.ImageMessage(model.RoleUser, "", "https://x/y.png")}, want: true},
|
|
{name: "image in later turn", msgs: []model.Message{
|
|
model.TextMessage(model.RoleUser, "hi"),
|
|
model.ImageMessage(model.RoleAssistant, "", "http://x/z.png"),
|
|
}, want: true},
|
|
}
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
if got := DetectMultimodal(tt.msgs); got != tt.want {
|
|
t.Errorf("DetectMultimodal() = %v, want %v", got, tt.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestParseImageRef(t *testing.T) {
|
|
t.Parallel()
|
|
tests := []struct {
|
|
name string
|
|
src string
|
|
wantSource ImageSource
|
|
}{
|
|
{name: "http url", src: "https://a/b.png", wantSource: SourceURL},
|
|
{name: "https url", src: "https://a/b.png", wantSource: SourceURL},
|
|
{name: "file uri", src: "file:///tmp/a.png", wantSource: SourcePath},
|
|
{name: "data uri", src: "data:image/png;base64,abc", wantSource: SourceBase64},
|
|
{name: "raw base64", src: "aGVsbG8=", wantSource: SourceBase64},
|
|
{name: "empty rejected", src: "", wantSource: ""},
|
|
}
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
ref, err := ParseImageRef(tt.src)
|
|
if tt.src == "" {
|
|
if err == nil {
|
|
t.Fatal("want error for empty source, got nil")
|
|
}
|
|
return
|
|
}
|
|
if err != nil {
|
|
t.Fatalf("ParseImageRef error: %v", err)
|
|
}
|
|
if ref.Source != tt.wantSource {
|
|
t.Errorf("Source = %q, want %q", ref.Source, tt.wantSource)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestResolve_Base64(t *testing.T) {
|
|
t.Parallel()
|
|
ref, err := ParseImageRef("data:image/png;base64,aGVsbG8=")
|
|
if err != nil {
|
|
t.Fatalf("ParseImageRef error: %v", err)
|
|
}
|
|
mime, data, err := Resolve(context.Background(), ref, nil)
|
|
if err != nil {
|
|
t.Fatalf("Resolve error: %v", err)
|
|
}
|
|
if mime != "image/png" {
|
|
t.Errorf("mime = %q, want image/png", mime)
|
|
}
|
|
if string(data) != "hello" {
|
|
t.Errorf("data = %q, want hello", data)
|
|
}
|
|
}
|
|
|
|
func TestResolve_DataURI(t *testing.T) {
|
|
t.Parallel()
|
|
got := DataURI("image/webp", []byte{1, 2, 3})
|
|
if !strings.HasPrefix(got, "data:image/webp;base64,") {
|
|
t.Errorf("DataURI = %q, want data:image/webp;base64,", got)
|
|
}
|
|
}
|