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

   1 package main
   2
   3 import (
   4         "encoding/json"
   5         "errors"
   6         "fmt"
   7         "io"
   8         "mime/multipart"
   9         "net/http"
  10         "net/http/httptest"
  11         "strings"
  12         "sync"
  13         "testing"
  14         "time"
  15 )
  16
  17 const (
  18         // telegramTestToken is the bot token of the tests. It has the form of a
  19         // real token.
  20         telegramTestToken = "123456:TEST-secret_token"
  21
  22         // telegramTestOwner is the user ID of the owner of the test bot.
  23         telegramTestOwner = 4242
  24 )
  25
  26 // TelegramFakeCall is one call of the fake Bot API. OK tells if the call got a
  27 // success answer.
  28 type TelegramFakeCall struct {
  29         method string
  30         params map[string]any
  31         ok     bool
  32 }
  33
  34 // TelegramFake is a fake Bot API server for telegramTestToken. getUpdates
  35 // returns the added updates from the offset on, or waits a short time when
  36 // there are none. The other methods succeed, unless fail gives an error answer
  37 // for a call. An answer with HTTP 429 asks for a wait of one second. When stall
  38 // is true for a call, the call gets no answer until the client gives it up or
  39 // the test ends. Fail and stall run with the mutex held.
  40 type TelegramFake struct {
  41         *httptest.Server
  42
  43         mutex   sync.Mutex
  44         updates []TelegramUpdate
  45         calls   []TelegramFakeCall
  46         fail    func(call TelegramFakeCall) (int, string)
  47         stall   func(call TelegramFakeCall) bool
  48         added   chan struct{}
  49         stopped chan struct{}
  50         images  map[string][]byte
  51 }
  52
  53 // newTelegramFake starts a fake Bot API server that stops at the end of the
  54 // test.
  55 func newTelegramFake(t *testing.T) *TelegramFake {
  56         t.Helper()
  57
  58         f := &TelegramFake{
  59                 added:   make(chan struct{}, 1),
  60                 stopped: make(chan struct{}),
  61         }
  62
  63         f.Server = httptest.NewServer(http.HandlerFunc(f.serve))
  64         t.Cleanup(f.Close)
  65         t.Cleanup(func() {
  66                 close(f.stopped)
  67         })
  68
  69         return f
  70 }
  71
  72 // newTelegramLocalFake returns the same fake API without network goroutines.
  73 // Its RoundTrip handles requests directly for tests with fake time.
  74 func newTelegramLocalFake(t *testing.T) *TelegramFake {
  75         t.Helper()
  76
  77         f := &TelegramFake{
  78                 Server:  &httptest.Server{URL: "https://example.com"},
  79                 added:   make(chan struct{}, 1),
  80                 stopped: make(chan struct{}),
  81         }
  82
  83         t.Cleanup(func() {
  84                 close(f.stopped)
  85         })
  86
  87         return f
  88 }
  89
  90 func (f *TelegramFake) serve(w http.ResponseWriter, r *http.Request) {
  91         id, file := strings.CutPrefix(r.URL.Path, "/file/bot"+telegramTestToken+"/")
  92         if file {
  93                 f.mutex.Lock()
  94
  95                 data, exists := f.images[id]
  96                 f.mutex.Unlock()
  97
  98                 if !exists {
  99                         http.NotFound(w, r)
 100
 101                         return
 102                 }
 103
 104                 // The client can cancel a media request before this reply is
 105                 // written.
 106                 w.Write(data) // nolint: errcheck
 107
 108                 return
 109         }
 110
 111         method, found := strings.CutPrefix(r.URL.Path, "/bot"+telegramTestToken+"/")
 112         if !found {
 113                 http.Error(w, `{"ok":false,"error_code":404,"description":"Not Found"}`, http.StatusNotFound)
 114
 115                 return
 116         }
 117
 118         var params map[string]any
 119
 120         var err error
 121
 122         if strings.HasPrefix(r.Header.Get("Content-Type"), "multipart/form-data;") {
 123                 err = r.ParseMultipartForm(imageMaxBytes + (1 << 20))
 124                 if err == nil {
 125                         params = make(map[string]any)
 126                         for key, values := range r.MultipartForm.Value {
 127                                 params[key] = values[0]
 128                         }
 129
 130                         for key, files := range r.MultipartForm.File {
 131                                 var file multipart.File
 132
 133                                 file, err = files[0].Open()
 134                                 if err != nil {
 135                                         break
 136                                 }
 137
 138                                 params[key], err = io.ReadAll(file)
 139                                 err = errors.Join(err, file.Close())
 140                                 params[key+"_name"] = files[0].Filename
 141                                 if err != nil {
 142                                         break
 143                                 }
 144                         }
 145
 146                         err = errors.Join(err, r.MultipartForm.RemoveAll())
 147                 }
 148         } else {
 149                 err = json.NewDecoder(r.Body).Decode(&params)
 150         }
 151
 152         if err != nil {
 153                 http.Error(w, err.Error(), http.StatusBadRequest)
 154
 155                 return
 156         }
 157
 158         call := TelegramFakeCall{
 159                 method: method,
 160                 params: params,
 161         }
 162
 163         f.mutex.Lock()
 164
 165         status, description := 0, ""
 166         if f.fail != nil {
 167                 status, description = f.fail(call)
 168         }
 169
 170         stall := f.stall != nil && f.stall(call)
 171         f.mutex.Unlock()
 172
 173         if stall {
 174                 f.record(call)
 175
 176                 select {
 177                 case <-r.Context().Done():
 178                 case <-f.stopped:
 179                 }
 180
 181                 return
 182         }
 183
 184         if status != 0 {
 185                 f.record(call)
 186                 w.WriteHeader(status)
 187
 188                 wait := 0
 189                 if status == http.StatusTooManyRequests {
 190                         wait = 1
 191                 }
 192
 193                 // A test can cancel the request after the handler selects its
 194                 // reply.
 195                 fmt.Fprintf(w, `{"ok":false,"error_code":%d,"description":%q,"parameters":{"retry_after":%d}}`, status, description, wait) // nolint: errcheck
 196
 197                 return
 198         }
 199
 200         call.ok = true
 201         f.record(call)
 202
 203         if method == "getFile" {
 204                 // A cancelled client does not need a response body.
 205                 json.NewEncoder(w).Encode(map[string]any{ // nolint: errcheck
 206                         "ok":     true,
 207                         "result": map[string]any{"file_path": params["file_id"]},
 208                 })
 209
 210                 return
 211         }
 212
 213         if method != "getUpdates" {
 214                 // A cancelled client does not need a response body.
 215                 io.WriteString(w, `{"ok":true,"result":true}`) // nolint: errcheck
 216
 217                 return
 218         }
 219
 220         offset := int64(params["offset"].(float64))
 221         updates := f.from(offset)
 222         if len(updates) == 0 {
 223                 select {
 224                 case <-f.added:
 225                 case <-time.After(50 * time.Millisecond):
 226                 case <-r.Context().Done():
 227                 }
 228
 229                 updates = f.from(offset)
 230         }
 231
 232         // A cancelled client does not need a response body.
 233         json.NewEncoder(w).Encode(map[string]any{ // nolint: errcheck
 234                 "ok":     true,
 235                 "result": updates,
 236         })
 237 }
 238
 239 // record keeps call.
 240 func (f *TelegramFake) record(call TelegramFakeCall) {
 241         f.mutex.Lock()
 242         defer f.mutex.Unlock()
 243
 244         f.calls = append(f.calls, call)
 245 }
 246
 247 // from returns the updates from offset on, and forgets the ones before it, as
 248 // Telegram does.
 249 func (f *TelegramFake) from(offset int64) []TelegramUpdate {
 250         f.mutex.Lock()
 251         defer f.mutex.Unlock()
 252
 253         var updates []TelegramUpdate
 254
 255         kept := f.updates[:0]
 256         for _, update := range f.updates {
 257                 if update.UpdateID >= offset {
 258                         updates = append(updates, update)
 259                         kept = append(kept, update)
 260                 }
 261         }
 262
 263         f.updates = kept
 264
 265         return updates
 266 }
 267
 268 // Add makes update available to getUpdates.
 269 func (f *TelegramFake) Add(update TelegramUpdate) {
 270         f.mutex.Lock()
 271
 272         f.updates = append(f.updates, update)
 273         f.mutex.Unlock()
 274
 275         select {
 276         case f.added <- struct{}{}:
 277         default:
 278         }
 279 }
 280
 281 // Sent returns the texts of the sendMessage calls and the Markdown of the
 282 // sendRichMessage calls that succeeded, in order.
 283 func (f *TelegramFake) Sent() []string {
 284         var texts []string
 285
 286         for _, call := range f.Calls() {
 287                 switch {
 288                 case !call.ok:
 289                 case call.method == "sendMessage":
 290                         texts = append(texts, call.params["text"].(string))
 291                 case call.method == "sendRichMessage":
 292                         texts = append(texts, call.params["rich_message"].(map[string]any)["markdown"].(string))
 293                 }
 294         }
 295
 296         return texts
 297 }
 298
 299 // Calls returns the calls so far.
 300 func (f *TelegramFake) Calls() []TelegramFakeCall {
 301         f.mutex.Lock()
 302         defer f.mutex.Unlock()
 303
 304         return append([]TelegramFakeCall(nil), f.calls...)
 305 }
 306
 307 func (f *TelegramFake) Image(id string, data []byte) {
 308         f.mutex.Lock()
 309         defer f.mutex.Unlock()
 310
 311         if f.images == nil {
 312                 f.images = make(map[string][]byte)
 313         }
 314
 315         f.images[id] = data
 316 }
 317
 318 // RoundTrip handles a request in memory, so synctest can control its timers.
 319 func (f *TelegramFake) RoundTrip(request *http.Request) (*http.Response, error) {
 320         response := httptest.NewRecorder()
 321         f.serve(response, request)
 322
 323         ctx := request.Context()
 324         if ctx.Err() != nil {
 325                 return nil, ctx.Err()
 326         }
 327
 328         return response.Result(), nil
 329 }
 330
 331 // telegramOwnerMessage returns a text message of the owner in the private chat.
 332 func telegramOwnerMessage(text string) *TelegramMessage {
 333         return &TelegramMessage{
 334                 Chat: TelegramChat{
 335                         ID:   telegramTestOwner,
 336                         Type: "private",
 337                 },
 338                 From: &TelegramUser{
 339                         ID: telegramTestOwner,
 340                 },
 341                 Text: text,
 342         }
 343 }