306 lines
12 KiB
Go
306 lines
12 KiB
Go
package server
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/base64"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"net/url"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
)
|
|
|
|
// fakeVertex is the image service: it records what it was asked and answers
|
|
// like the real one (a job to poll, a picture when done).
|
|
type fakeVertex struct {
|
|
srv *httptest.Server
|
|
mu sync.Mutex
|
|
auth []string
|
|
gen map[string]any
|
|
edit struct {
|
|
prompt string
|
|
image []byte
|
|
name string
|
|
}
|
|
polls int
|
|
failAs string // an error_code the job ends with
|
|
}
|
|
|
|
func newFakeVertex(t *testing.T) *fakeVertex {
|
|
f := &fakeVertex{}
|
|
f.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
f.auth = append(f.auth, r.Header.Get("Authorization"))
|
|
// The pictures are open by their address, as at the real service.
|
|
if r.URL.Path == "/uploads/userphoto/ai/t.png" {
|
|
w.Header().Set("Content-Type", "image/png")
|
|
_, _ = w.Write([]byte("\x89PNG-original"))
|
|
return
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
if r.Header.Get("Authorization") != "Bearer vx_good" {
|
|
w.WriteHeader(401)
|
|
fmt.Fprint(w, `{"ok":false,"error_code":"unauthorized","error":"Неверный ключ"}`)
|
|
return
|
|
}
|
|
switch {
|
|
case r.URL.Path == "/api/vertex/v1/generate":
|
|
_ = json.NewDecoder(r.Body).Decode(&f.gen)
|
|
w.WriteHeader(202)
|
|
fmt.Fprint(w, `{"ok":true,"job_id":"job_gen1","status":"queued","cost":5}`)
|
|
case r.URL.Path == "/api/vertex/v1/edit":
|
|
if err := r.ParseMultipartForm(32 << 20); err != nil {
|
|
w.WriteHeader(400)
|
|
fmt.Fprint(w, `{"ok":false,"error_code":"invalid_image","error":"нет фото"}`)
|
|
return
|
|
}
|
|
f.edit.prompt = r.FormValue("prompt")
|
|
file, hdr, err := r.FormFile("image")
|
|
if err == nil {
|
|
f.edit.image, _ = io.ReadAll(file)
|
|
f.edit.name = hdr.Filename
|
|
}
|
|
w.WriteHeader(202)
|
|
fmt.Fprint(w, `{"ok":true,"job_id":"job_edit1","status":"queued","cost":2}`)
|
|
case strings.HasPrefix(r.URL.Path, "/api/vertex/v1/jobs/"):
|
|
f.polls++
|
|
if f.failAs != "" {
|
|
fmt.Fprintf(w, `{"ok":true,"job_id":"x","status":"failed","error_code":%q,"error":"сбой"}`, f.failAs)
|
|
return
|
|
}
|
|
if f.polls < 2 {
|
|
fmt.Fprint(w, `{"ok":true,"job_id":"x","status":"running"}`)
|
|
return
|
|
}
|
|
fmt.Fprintf(w, `{"ok":true,"job_id":"x","status":"done","photo_url":"%s/uploads/userphoto/ai/t.png","cost":5,"balance":42}`, f.srv.URL)
|
|
default:
|
|
http.NotFound(w, r)
|
|
}
|
|
}))
|
|
t.Cleanup(f.srv.Close)
|
|
return f
|
|
}
|
|
|
|
// setImage points the panel at the fake service with a key.
|
|
func setImage(t *testing.T, s *Server, f *fakeVertex, key string) {
|
|
t.Helper()
|
|
_ = s.db.SetConfig(cfgImageURL, f.srv.URL)
|
|
_ = s.db.SetConfig(cfgImageKey, key)
|
|
}
|
|
|
|
func apiCall(t *testing.T, s *Server, c *http.Cookie, method, path, body string) string {
|
|
t.Helper()
|
|
req := httptest.NewRequest(method, path, strings.NewReader(body))
|
|
req.AddCookie(c)
|
|
w := httptest.NewRecorder()
|
|
s.Handler().ServeHTTP(w, req)
|
|
return w.Body.String()
|
|
}
|
|
|
|
// A picture from nothing: the prompt, model and ratio go to the service with
|
|
// the key; the job's progress and its picture come back, the key never.
|
|
func TestImageGenerate(t *testing.T) {
|
|
s, c := newTestServer(t)
|
|
f := newFakeVertex(t)
|
|
setImage(t, s, f, "vx_good")
|
|
got := apiCall(t, s, c, "POST", "/api/panel/ai/image", `{"mode":"generate","prompt":"cat in space","model":"artist","ratio":"3:4"}`)
|
|
if !strings.Contains(got, `"ok":true`) || !strings.Contains(got, `"job_id":"job_gen1"`) || strings.Contains(got, "vx_good") {
|
|
t.Fatalf("create: %s", got)
|
|
}
|
|
if f.gen["prompt"] != "cat in space" || f.gen["model"] != "artist" || f.gen["ratio"] != "3:4" || f.gen["enhance"] != true {
|
|
t.Fatalf("what the service got: %v", f.gen)
|
|
}
|
|
if got := apiCall(t, s, c, "GET", "/api/panel/ai/image/job_gen1?wait=1", ""); !strings.Contains(got, `"status":"running"`) {
|
|
t.Fatalf("first poll: %s", got)
|
|
}
|
|
got = apiCall(t, s, c, "GET", "/api/panel/ai/image/job_gen1?wait=1", "")
|
|
if !strings.Contains(got, `"status":"done"`) || !strings.Contains(got, "/uploads/userphoto/ai/t.png") || !strings.Contains(got, `"balance":42`) || strings.Contains(got, "vx_good") {
|
|
t.Fatalf("done: %s", got)
|
|
}
|
|
}
|
|
|
|
// Unknown models and ratios are not sent on: the defaults go instead.
|
|
func TestImageGenerateDefaults(t *testing.T) {
|
|
s, c := newTestServer(t)
|
|
f := newFakeVertex(t)
|
|
setImage(t, s, f, "vx_good")
|
|
apiCall(t, s, c, "POST", "/api/panel/ai/image", `{"mode":"generate","prompt":"x","model":"gpt","ratio":"7:1"}`)
|
|
if f.gen["model"] != "fast" || f.gen["ratio"] != "1:1" {
|
|
t.Fatalf("defaults: %v", f.gen)
|
|
}
|
|
}
|
|
|
|
// Editing sends the user's photo as a file with the instruction.
|
|
func TestImageEditFromAttachedPhoto(t *testing.T) {
|
|
s, c := newTestServer(t)
|
|
f := newFakeVertex(t)
|
|
setImage(t, s, f, "vx_good")
|
|
photo := []byte("\xff\xd8\xffJPEGDATA")
|
|
b, _ := json.Marshal(map[string]any{"mode": "edit", "prompt": "сделай закат",
|
|
"image": map[string]string{"mime": "image/jpeg", "data": base64.StdEncoding.EncodeToString(photo)}})
|
|
got := apiCall(t, s, c, "POST", "/api/panel/ai/image", string(b))
|
|
if !strings.Contains(got, `"job_id":"job_edit1"`) {
|
|
t.Fatalf("edit: %s", got)
|
|
}
|
|
if f.edit.prompt != "сделай закат" || !bytes.Equal(f.edit.image, photo) || !strings.HasSuffix(f.edit.name, ".jpg") {
|
|
t.Fatalf("what the service got: %q %q %q", f.edit.prompt, f.edit.image, f.edit.name)
|
|
}
|
|
// A format the service does not take is told so, not sent.
|
|
b, _ = json.Marshal(map[string]any{"mode": "edit", "prompt": "x", "image": map[string]string{"mime": "image/gif", "data": "R0lG"}})
|
|
if got := apiCall(t, s, c, "POST", "/api/panel/ai/image", string(b)); !strings.Contains(got, `"ok":false`) || !strings.Contains(got, "png") {
|
|
t.Fatalf("gif: %s", got)
|
|
}
|
|
}
|
|
|
|
// A picture made before can be edited again: fetched from the service's own
|
|
// address only - never from another.
|
|
func TestImageEditFromEarlierPicture(t *testing.T) {
|
|
s, c := newTestServer(t)
|
|
f := newFakeVertex(t)
|
|
setImage(t, s, f, "vx_good")
|
|
b, _ := json.Marshal(map[string]any{"mode": "edit", "prompt": "темнее", "image_url": f.srv.URL + "/uploads/userphoto/ai/t.png"})
|
|
if got := apiCall(t, s, c, "POST", "/api/panel/ai/image", string(b)); !strings.Contains(got, `"ok":true`) {
|
|
t.Fatalf("by url: %s", got)
|
|
}
|
|
if !bytes.Equal(f.edit.image, []byte("\x89PNG-original")) || !strings.HasSuffix(f.edit.name, ".png") {
|
|
t.Fatalf("the picture sent: %q %q", f.edit.image, f.edit.name)
|
|
}
|
|
b, _ = json.Marshal(map[string]any{"mode": "edit", "prompt": "x", "image_url": "http://127.0.0.1:9/secret.png"})
|
|
if got := apiCall(t, s, c, "POST", "/api/panel/ai/image", string(b)); !strings.Contains(got, `"ok":false`) {
|
|
t.Fatalf("a foreign address was fetched: %s", got)
|
|
}
|
|
}
|
|
|
|
// What goes wrong is told in words: no key, a wrong key, the service's own
|
|
// codes, a job that failed.
|
|
func TestImageErrors(t *testing.T) {
|
|
s, c := newTestServer(t)
|
|
f := newFakeVertex(t)
|
|
if got := apiCall(t, s, c, "POST", "/api/panel/ai/image", `{"mode":"generate","prompt":"x"}`); !strings.Contains(got, "ключ") {
|
|
t.Fatalf("no key: %s", got)
|
|
}
|
|
setImage(t, s, f, "vx_bad")
|
|
if got := apiCall(t, s, c, "POST", "/api/panel/ai/image", `{"mode":"generate","prompt":"x"}`); !strings.Contains(got, `"ok":false`) || !strings.Contains(got, "ключ") {
|
|
t.Fatalf("wrong key: %s", got)
|
|
}
|
|
setImage(t, s, f, "vx_good")
|
|
f.failAs = "insufficient_credits"
|
|
if got := apiCall(t, s, c, "GET", "/api/panel/ai/image/job_gen1", ""); !strings.Contains(got, `"status":"failed"`) || !strings.Contains(got, "кредит") {
|
|
t.Fatalf("failed job: %s", got)
|
|
}
|
|
if got := apiCall(t, s, c, "GET", "/api/panel/ai/image/..%2Fme", ""); !strings.Contains(got, `"ok":false`) {
|
|
t.Fatalf("a job id out of shape: %s", got)
|
|
}
|
|
if got := apiCall(t, s, c, "POST", "/api/panel/ai/image", `{"mode":"generate","prompt":""}`); !strings.Contains(got, `"ok":false`) {
|
|
t.Fatalf("empty prompt: %s", got)
|
|
}
|
|
}
|
|
|
|
// The assistant off in the profile means no pictures either.
|
|
func TestImageNeedsAssistantOn(t *testing.T) {
|
|
s, c := newTestServer(t)
|
|
f := newFakeVertex(t)
|
|
setImage(t, s, f, "vx_good")
|
|
a, _ := s.db.AdminByUsername("adm")
|
|
_ = s.db.SetConfig(aiUserKey(a.ID), "0")
|
|
if got := apiCall(t, s, c, "POST", "/api/panel/ai/image", `{"mode":"generate","prompt":"x"}`); !strings.Contains(got, "выключен") {
|
|
t.Fatalf("off: %s", got)
|
|
}
|
|
}
|
|
|
|
// The key is saved from Настройки, shown back to the admin, and the chat
|
|
// learns only that pictures are on.
|
|
func TestImageKeySetting(t *testing.T) {
|
|
s, c := newTestServer(t)
|
|
postAIModels(t, s, c, [][7]string{{"", "M", "gradio", "", "", "", ""}})
|
|
f := url.Values{"csrf": {"tok"}, "m_type": {"gradio"}, "ai_image_key": {"vx_secret"}}
|
|
req := httptest.NewRequest(http.MethodPost, "/settings/ai", strings.NewReader(f.Encode()))
|
|
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
|
req.AddCookie(c)
|
|
req.AddCookie(&http.Cookie{Name: csrfCookie, Value: "tok"})
|
|
s.Handler().ServeHTTP(httptest.NewRecorder(), req)
|
|
if _, k := s.imageConf(); k != "vx_secret" {
|
|
t.Fatalf("key not saved: %q", k)
|
|
}
|
|
if body := get(t, s, c, "/settings?frag=1").Body.String(); !strings.Contains(body, `value="vx_secret"`) {
|
|
t.Fatal("the settings page does not show the key")
|
|
}
|
|
state := get(t, s, c, "/api/panel/ai").Body.String()
|
|
if strings.Contains(state, "vx_secret") || !strings.Contains(state, `"image":true`) {
|
|
t.Fatalf("state: %s", state)
|
|
}
|
|
}
|
|
|
|
// With a key the model is told it can draw; without, it is not.
|
|
func TestImagePromptOnlyWithKey(t *testing.T) {
|
|
s, _ := newTestServer(t)
|
|
if strings.Contains(s.imagePrompt(), "```image") {
|
|
t.Fatal("told about pictures with no key")
|
|
}
|
|
_ = s.db.SetConfig(cfgImageKey, "vx_good")
|
|
p := s.imagePrompt()
|
|
for _, want := range []string{"```image\n", "```image edit", "fast", "artist", "0.2plus", "16:9"} {
|
|
if !strings.Contains(p, want) {
|
|
t.Fatalf("the prompt lacks %q:\n%s", want, p)
|
|
}
|
|
}
|
|
}
|
|
|
|
// The chat gives the model the image instructions when pictures are on.
|
|
func TestAIChatMentionsImageTool(t *testing.T) {
|
|
s, c := newTestServer(t)
|
|
var system string
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
var b map[string]any
|
|
_ = json.NewDecoder(r.Body).Decode(&b)
|
|
system, _ = b["messages"].([]any)[0].(map[string]any)["content"].(string)
|
|
fmt.Fprint(w, openAIChunk("content", "ok")+"data: [DONE]\n\n")
|
|
}))
|
|
defer srv.Close()
|
|
postAIModels(t, s, c, [][7]string{{"", "M", "openai", srv.URL, "m", "", ""}})
|
|
ask := func() {
|
|
apiCall(t, s, c, "POST", "/api/panel/ai/chat", `{"message":"нарисуй кота","context":"x"}`)
|
|
}
|
|
ask()
|
|
if strings.Contains(system, "```image") {
|
|
t.Fatal("image tool offered with no key")
|
|
}
|
|
_ = s.db.SetConfig(cfgImageKey, "vx_good")
|
|
ask()
|
|
if !strings.Contains(system, "```image") {
|
|
t.Fatalf("image tool not offered with a key:\n%s", system)
|
|
}
|
|
}
|
|
|
|
// The description is improved by the service unless the model asks for it
|
|
// as it is.
|
|
func TestImageEnhanceSwitch(t *testing.T) {
|
|
s, c := newTestServer(t)
|
|
f := newFakeVertex(t)
|
|
setImage(t, s, f, "vx_good")
|
|
apiCall(t, s, c, "POST", "/api/panel/ai/image", `{"mode":"generate","prompt":"x"}`)
|
|
if f.gen["enhance"] != true {
|
|
t.Fatalf("by default: %v", f.gen)
|
|
}
|
|
apiCall(t, s, c, "POST", "/api/panel/ai/image", `{"mode":"generate","prompt":"x","enhance":false}`)
|
|
if f.gen["enhance"] != false {
|
|
t.Fatalf("switched off: %v", f.gen)
|
|
}
|
|
apiCall(t, s, c, "POST", "/api/panel/ai/image", `{"mode":"generate","prompt":"x","enhance":true}`)
|
|
if f.gen["enhance"] != true {
|
|
t.Fatalf("switched on: %v", f.gen)
|
|
}
|
|
{
|
|
p := s.imagePrompt()
|
|
_ = s.db.SetConfig(cfgImageKey, "vx_good")
|
|
if !strings.Contains(s.imagePrompt(), "noenhance") {
|
|
t.Fatalf("the model is not told of the switch:\n%s", p)
|
|
}
|
|
}
|
|
}
|