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