git.giacomolli.uk - onechat/blob - image_show_test.go

   1 package main
   2
   3 import (
   4         "encoding/base64"
   5         "encoding/json"
   6         "net/http"
   7         "path/filepath"
   8         "testing"
   9
  10         "github.com/stretchr/testify/require"
  11 )
  12
  13 func TestImageShowUsesStoredImageForTelegramAndTerminal(t *testing.T) {
  14         picture := imageTestPNG(t, 5, 4)
  15         store := &ImageStore{dir: filepath.Join(t.TempDir(), imageDirName)}
  16         name, err := store.Save(picture)
  17         require.NoError(t, err)
  18
  19         telegram := newTelegramFake(t)
  20         bot, logs := telegramStartTestBot(t, telegram)
  21         box := &Toolbox{imageShow: &ImageShow{
  22                 store: store,
  23                 bot:   bot,
  24         }}
  25         args, err := json.Marshal(ToolImageShowArguments{
  26                 Image:   name,
  27                 Caption: "A red point.",
  28         })
  29         require.NoError(t, err)
  30
  31         output := box.run(t.Context(), "image_show", string(args))
  32         require.Equal(t, imageText([]string{name}, "A red point.")+"\n"+filepath.Join(store.dir, name), output)
  33
  34         calls := telegram.Calls()
  35         require.Len(t, calls, 1)
  36         require.Equal(t, "sendPhoto", calls[0].method)
  37         require.Equal(t, "4242", calls[0].params["chat_id"])
  38         require.Equal(t, "A red point.", calls[0].params["caption"])
  39         require.Equal(t, name, calls[0].params["photo_name"])
  40         require.Equal(t, picture, calls[0].params["photo"])
  41         require.Equal(t, "<6>image sent\n", logs.String())
  42 }
  43
  44 func TestImageShowRefusesUnknownNamesBeforeSending(t *testing.T) {
  45         telegram := newTelegramFake(t)
  46         bot, _ := telegramStartTestBot(t, telegram)
  47         show := &ImageShow{
  48                 store: &ImageStore{dir: t.TempDir()},
  49                 bot:   bot,
  50         }
  51
  52         for _, name := range []string{"../image.png", "unknown", "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa.png"} {
  53                 _, err := show.Show(t.Context(), name, "Do not send this.")
  54                 require.Error(t, err)
  55         }
  56
  57         require.Empty(t, telegram.Calls())
  58
  59         box := &Toolbox{}
  60         require.Equal(t, `error: unknown tool "image_show"`, box.run(t.Context(), "image_show", `{"image":"unknown","caption":""}`))
  61 }
  62
  63 func TestImageShowFallsBackOnlyWhenTelegramRefusesPhoto(t *testing.T) {
  64         testCases := []struct {
  65                 name    string
  66                 status  int
  67                 methods []string
  68                 failure bool
  69         }{
  70                 {
  71                         name:    "refused photo sends document",
  72                         status:  http.StatusBadRequest,
  73                         methods: []string{"sendPhoto", "sendDocument"},
  74                 },
  75                 {
  76                         name:    "uncertain server error does not send twice",
  77                         status:  http.StatusBadGateway,
  78                         methods: []string{"sendPhoto"},
  79                         failure: true,
  80                 },
  81         }
  82         for _, tc := range testCases {
  83                 t.Run(tc.name, func(t *testing.T) {
  84                         telegram := newTelegramFake(t)
  85                         telegram.fail = func(call TelegramFakeCall) (int, string) {
  86                                 if call.method == "sendPhoto" {
  87                                         return tc.status, "photo refused"
  88                                 }
  89
  90                                 return 0, ""
  91                         }
  92                         store := &ImageStore{dir: t.TempDir()}
  93                         picture := imageTestPNG(t, 3, 3)
  94                         name, err := store.Save(picture)
  95                         require.NoError(t, err)
  96
  97                         bot, _ := telegramStartTestBot(t, telegram)
  98                         show := &ImageShow{
  99                                 store: store,
 100                                 bot:   bot,
 101                         }
 102                         _, err = show.Show(t.Context(), name, "Caption")
 103                         if tc.failure {
 104                                 require.Error(t, err)
 105                         } else {
 106                                 require.NoError(t, err)
 107                         }
 108
 109                         var methods []string
 110
 111                         for _, call := range telegram.Calls() {
 112                                 methods = append(methods, call.method)
 113                                 if call.method == "sendDocument" {
 114                                         require.Equal(t, picture, call.params["document"])
 115                                         require.Equal(t, name, call.params["document_name"])
 116                                         require.Equal(t, "Caption", call.params["caption"])
 117                                 }
 118                         }
 119
 120                         require.Equal(t, tc.methods, methods)
 121                 })
 122         }
 123 }
 124
 125 func TestAgentShowsImageInTerminalWithoutTelegram(t *testing.T) {
 126         server := newModelFake(t)
 127         test := newTestTurnLoop(t, server)
 128         name, err := test.agent.Images().Save(imageTestPNG(t, 2, 2))
 129         require.NoError(t, err)
 130
 131         output := test.loop.tools.run(t.Context(), "image_show", `{"image":"`+name+`","caption":"Here it is."}`)
 132         require.Equal(t, imageText([]string{name}, "Here it is.")+"\n"+filepath.Join(test.agent.Images().dir, name), output)
 133         require.Contains(t, terminalEventText(Event{
 134                 Kind: "echo",
 135                 Text: output,
 136         }), filepath.Join(test.agent.Images().dir, name))
 137 }
 138
 139 func TestTellImageMarkerGivesTheChildPixels(t *testing.T) {
 140         server := newModelFake(t)
 141         test := newDaemonMessageTest(t, t.TempDir(), server)
 142         picture := imageTestPNG(t, 4, 2)
 143         name, err := test.agents[0].agent.Images().Save(picture)
 144         require.NoError(t, err)
 145         test.Start(t, 0, 1)
 146
 147         arguments, err := json.Marshal(ToolTellArguments{
 148                 Name:    "Garden",
 149                 Message: imageText([]string{name}, "Review this plan."),
 150         })
 151         require.NoError(t, err)
 152         require.Contains(t, test.loops[0].tools.run(t.Context(), "tell", string(arguments)), "Sent.")
 153         test.events[1].WaitFor(t, "idle", 1)
 154
 155         var childCall *ModelFakeCall
 156
 157         for i := range server.Calls() {
 158                 call := server.Call(i)
 159                 if call.body["model"] == "gpt-test-garden" {
 160                         childCall = &call
 161                         break
 162                 }
 163         }
 164
 165         require.NotNil(t, childCall)
 166         require.Contains(t, string(childCall.raw), "input_image")
 167         require.Contains(t, string(childCall.raw), base64.StdEncoding.EncodeToString(picture))
 168         require.Contains(t, test.events[1].Last("user"), "[image "+name+"]")
 169         require.NotContains(t, promptToolNames(test.agents[1].agent.Tools()), "image_show")
 170 }
 171
 172 func TestAgentShowsTelegramReceivedImageToItsOwner(t *testing.T) {
 173         picture := imageTestPNG(t, 6, 4)
 174         store := &ImageStore{dir: t.TempDir()}
 175         name, err := store.Save(picture)
 176         require.NoError(t, err)
 177
 178         telegram := newTelegramFake(t)
 179         telegram.Image("owner-photo", picture)
 180
 181         bot, logs := telegramStartTestBot(t, telegram)
 182         server := newModelFake(t)
 183         server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
 184                 if call.number == 0 {
 185                         arguments, err := json.Marshal(ToolImageShowArguments{
 186                                 Image:   name,
 187                                 Caption: "Here is your photo.",
 188                         })
 189                         require.NoError(t, err)
 190                         modelWriteEvents(w, modelCompletedEvent([]any{modelFunctionCallItem("fc_1", "call_1", "image_show", string(arguments))}, 100, 0, 5))
 191
 192                         return
 193                 }
 194
 195                 modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("Shown.")}, 100, 0, 5))
 196         }
 197         socket := daemonStartTestWithBot(t, server, bot)
 198         terminal := terminalStartTestClient(t, socket)
 199         terminal.WaitFor(t, "[connected to OneChat]")
 200
 201         message := telegramOwnerMessage("")
 202         message.Caption = "Show this photo."
 203         message.Photo = []TelegramPhoto{{
 204                 FileID: "owner-photo",
 205                 Width:  6,
 206                 Height: 4,
 207         }}
 208         telegram.Add(TelegramUpdate{
 209                 UpdateID: 1,
 210                 Message:  message,
 211         })
 212         terminal.WaitFor(t, "Shown.\n[done]")
 213         require.Contains(t, terminal.Output(), "[result] [image "+name+"]")
 214         require.Contains(t, terminal.Output(), filepath.Join(imageDirName, name))
 215
 216         var photos []TelegramFakeCall
 217
 218         for _, call := range telegram.Calls() {
 219                 if call.method == "sendPhoto" {
 220                         photos = append(photos, call)
 221                 }
 222         }
 223
 224         require.Len(t, photos, 1)
 225         require.Equal(t, picture, photos[0].params["photo"])
 226         require.Equal(t, "4242", photos[0].params["chat_id"])
 227         require.Equal(t, "Here is your photo.", photos[0].params["caption"])
 228         require.NotContains(t, logs.String(), "Here is your photo.")
 229         require.NotContains(t, logs.String(), "Show this photo.")
 230         require.NotContains(t, logs.String(), telegramTestToken)
 231 }