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

   1 package main
   2
   3 import (
   4         "bytes"
   5         "image"
   6         "image/jpeg"
   7         "net/http"
   8         "os"
   9         "path/filepath"
  10         "strings"
  11         "testing"
  12         "time"
  13
  14         "github.com/stretchr/testify/require"
  15 )
  16
  17 func TestTelegramSelectsLargestPhotoAtMost1568Pixels(t *testing.T) {
  18         server := newModelFake(t)
  19         test := newTestTurnLoop(t, server)
  20         telegram := newTelegramFake(t)
  21         bot, logs := telegramStartTestBot(t, telegram)
  22         picture := imageTestPNG(t, 8, 6)
  23         telegram.Image("selected", picture)
  24
  25         message := telegramOwnerMessage("")
  26         message.Caption = "Read this label.\n"
  27         message.Photo = []TelegramPhoto{
  28                 {
  29                         FileID: "small",
  30                         Width:  100,
  31                         Height: 200,
  32                 },
  33                 {
  34                         FileID: "too-large",
  35                         Width:  1600,
  36                         Height: 1200,
  37                 },
  38                 {
  39                         FileID: "selected",
  40                         Width:  1568,
  41                         Height: 1176,
  42                 },
  43                 {
  44                         FileID: "medium",
  45                         Width:  800,
  46                         Height: 600,
  47                 },
  48         }
  49         require.NoError(t, bot.handle(t.Context(), test.loop, nil, TelegramUpdate{
  50                 UpdateID: 1,
  51                 Message:  message,
  52         }))
  53
  54         test.events.WaitFor(t, "idle", 1)
  55
  56         files, err := os.ReadDir(test.agent.Images().dir)
  57         require.NoError(t, err)
  58         require.Len(t, files, 1)
  59
  60         logged, err := test.memory.Message(0)
  61         require.NoError(t, err)
  62         require.Equal(t, imageText([]string{files[0].Name()}, message.Caption), logged.Text)
  63         require.True(t, strings.HasSuffix(logged.Text, "\n"))
  64         require.Contains(t, string(server.Call(0).raw), "input_image")
  65         require.NotContains(t, logs.String(), "base64")
  66         require.NotContains(t, logs.String(), "Read this label")
  67
  68         for _, call := range telegram.Calls() {
  69                 if call.method == "getFile" {
  70                         require.Equal(t, "selected", call.params["file_id"])
  71                 }
  72         }
  73 }
  74
  75 func TestTelegramImageDocumentsUseContentFormatAndRejectOtherMedia(t *testing.T) {
  76         telegram := newTelegramFake(t)
  77         png := imageTestPNG(t, 2, 2)
  78
  79         var jpegBytes bytes.Buffer
  80
  81         require.NoError(t, jpeg.Encode(&jpegBytes, image.NewRGBA(image.Rect(0, 0, 2, 2)), nil))
  82
  83         // A made-up one-pixel lossless WebP image.
  84         webp := []byte{'R', 'I', 'F', 'F', 22, 0, 0, 0, 'W', 'E', 'B', 'P', 'V', 'P', '8', 'L', 10, 0, 0, 0, 0x2f, 0, 0, 0, 0, 0x07, 0x10, 0x11, 0x11, 0x88}
  85
  86         telegram.Image("png", png)
  87         telegram.Image("jpeg", jpegBytes.Bytes())
  88         telegram.Image("webp", webp)
  89         telegram.Image("invalid", []byte("not an image"))
  90         telegram.Image("large", bytes.Repeat([]byte{0}, imageMaxBytes+1))
  91
  92         api := newTelegramAPI(telegram.URL, telegramTestToken)
  93
  94         store := &ImageStore{dir: filepath.Join(t.TempDir(), imageDirName)}
  95         testCases := []struct {
  96                 name     string
  97                 document TelegramDocument
  98                 suffix   string
  99                 failure  string
 100         }{
 101                 {
 102                         name: "accepts PNG documents",
 103                         document: TelegramDocument{
 104                                 FileID:   "png",
 105                                 MimeType: "image/png",
 106                         },
 107                         suffix: ".png]",
 108                 },
 109                 {
 110                         name: "accepts JPEG documents",
 111                         document: TelegramDocument{
 112                                 FileID:   "jpeg",
 113                                 MimeType: "image/jpeg",
 114                         },
 115                         suffix: ".jpg]",
 116                 },
 117                 {
 118                         name: "accepts WebP documents",
 119                         document: TelegramDocument{
 120                                 FileID:   "webp",
 121                                 MimeType: "image/webp",
 122                         },
 123                         suffix: ".webp]",
 124                 },
 125                 {
 126                         name: "other documents keep text answer",
 127                         document: TelegramDocument{
 128                                 FileID:   "pdf",
 129                                 MimeType: "application/pdf",
 130                         },
 131                         failure: "only text messages are supported",
 132                 },
 133                 {
 134                         name: "other images refused",
 135                         document: TelegramDocument{
 136                                 FileID:   "gif",
 137                                 MimeType: "image/gif",
 138                         },
 139                         failure: "JPEG, PNG, or WebP",
 140                 },
 141                 {
 142                         name: "invalid image refused",
 143                         document: TelegramDocument{
 144                                 FileID:   "invalid",
 145                                 MimeType: "image/png",
 146                         },
 147                         failure: "JPEG, PNG, or WebP",
 148                 },
 149                 {
 150                         name: "large download refused",
 151                         document: TelegramDocument{
 152                                 FileID:   "large",
 153                                 MimeType: "image/png",
 154                         },
 155                         failure: "exceeds 20 MB",
 156                 },
 157                 {
 158                         name: "large metadata refused",
 159                         document: TelegramDocument{
 160                                 FileID:   "not-downloaded",
 161                                 MimeType: "image/png",
 162                                 FileSize: imageMaxBytes + 1,
 163                         },
 164                         failure: "exceeds 20 MB",
 165                 },
 166         }
 167
 168         for _, tc := range testCases {
 169                 t.Run(tc.name, func(t *testing.T) {
 170                         text, err := telegramImageText(t.Context(), api, store, &TelegramMessage{Document: &tc.document})
 171                         if tc.failure != "" {
 172                                 require.ErrorContains(t, err, tc.failure)
 173
 174                                 return
 175                         }
 176
 177                         require.NoError(t, err)
 178                         require.True(t, strings.HasSuffix(text, tc.suffix))
 179                 })
 180         }
 181 }
 182
 183 func TestTelegramRetriesImageDownloadBeforeAdvancingOffset(t *testing.T) {
 184         telegram := newTelegramFake(t)
 185         telegram.Image("photo", imageTestPNG(t, 2, 2))
 186
 187         tries := 0
 188         telegram.fail = func(call TelegramFakeCall) (int, string) {
 189                 if call.method == "getFile" {
 190                         tries++
 191                         if tries == 1 {
 192                                 return http.StatusServiceUnavailable, "try again"
 193                         }
 194                 }
 195
 196                 return 0, ""
 197         }
 198
 199         bot, logs := telegramStartTestBot(t, telegram)
 200         server := newModelFake(t)
 201         socket := daemonStartTestWithBot(t, server, bot)
 202         terminal := terminalStartTestClient(t, socket)
 203         terminal.WaitFor(t, "[connected to OneChat]")
 204
 205         message := telegramOwnerMessage("")
 206         message.Photo = []TelegramPhoto{{
 207                 FileID: "photo",
 208                 Width:  2,
 209                 Height: 2,
 210         }}
 211
 212         telegram.Add(TelegramUpdate{
 213                 UpdateID: 7,
 214                 Message:  message,
 215         })
 216
 217         terminal.WaitFor(t, "[done]")
 218         require.Eventually(t, func() bool {
 219                 offset, err := telegramLoadOffset(bot.offsetPath)
 220                 return err == nil && offset == 8
 221         }, 10*time.Second, time.Millisecond)
 222
 223         require.Contains(t, logs.String(), "receive failed")
 224         require.Contains(t, logs.String(), "receive back")
 225         require.NotContains(t, logs.String(), telegramTestToken)
 226         require.NotContains(t, logs.String(), "base64")
 227 }
 228
 229 func TestTelegramRefusesPhotoWhenEverySizeIsTooLarge(t *testing.T) {
 230         telegram := newTelegramFake(t)
 231         _, err := telegramImageText(t.Context(), newTelegramAPI(telegram.URL, telegramTestToken), &ImageStore{dir: t.TempDir()}, &TelegramMessage{Photo: []TelegramPhoto{{
 232                 FileID: "too-big",
 233                 Width:  1600,
 234                 Height: 1000,
 235         }}})
 236         require.ErrorContains(t, err, "1568 pixels")
 237         require.Empty(t, telegram.Calls())
 238 }
 239
 240 func TestTelegramStorageFailureDoesNotAcceptTheImage(t *testing.T) {
 241         telegram := newTelegramFake(t)
 242         telegram.Image("photo", imageTestPNG(t, 2, 2))
 243
 244         test := newTestTurnLoop(t, newModelFake(t))
 245         require.NoError(t, os.WriteFile(test.agent.Images().dir, []byte("blocked"), 0o600))
 246
 247         bot, _ := telegramStartTestBot(t, telegram)
 248         message := telegramOwnerMessage("")
 249         message.Photo = []TelegramPhoto{{
 250                 FileID: "photo",
 251                 Width:  2,
 252                 Height: 2,
 253         }}
 254
 255         err := bot.handle(t.Context(), test.loop, nil, TelegramUpdate{
 256                 UpdateID: 1,
 257                 Message:  message,
 258         })
 259         require.ErrorContains(t, err, "create image directory")
 260         require.Equal(t, 0, test.memory.Len())
 261
 262         _, found := bot.next()
 263         require.False(t, found)
 264 }