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 // Add makes update available to getUpdates.
  54 func (f *TelegramFake) Add(update TelegramUpdate) {
  55         f.mutex.Lock()
  56
  57         f.updates = append(f.updates, update)
  58         f.mutex.Unlock()
  59
  60         select {
  61         case f.added <- struct{}{}:
  62         default:
  63         }
  64 }
  65
  66 // Calls returns the calls so far.
  67 func (f *TelegramFake) Calls() []TelegramFakeCall {
  68         f.mutex.Lock()
  69         defer f.mutex.Unlock()
  70
  71         return append([]TelegramFakeCall(nil), f.calls...)
  72 }
  73
  74 // Sent returns the texts of the sendMessage calls and the Markdown of the
  75 // sendRichMessage calls that succeeded, in order.
  76 func (f *TelegramFake) Sent() []string {
  77         var texts []string
  78
  79         for _, call := range f.Calls() {
  80                 switch {
  81                 case !call.ok:
  82                 case call.method == "sendMessage":
  83                         texts = append(texts, call.params["text"].(string))
  84                 case call.method == "sendRichMessage":
  85                         texts = append(texts, call.params["rich_message"].(map[string]any)["markdown"].(string))
  86                 }
  87         }
  88
  89         return texts
  90 }
  91
  92 func (f *TelegramFake) Image(id string, data []byte) {
  93         f.mutex.Lock()
  94         defer f.mutex.Unlock()
  95
  96         if f.images == nil {
  97                 f.images = make(map[string][]byte)
  98         }
  99
 100         f.images[id] = data
 101 }
 102
 103 // RoundTrip handles a request in memory, so synctest can control its timers.
 104 func (f *TelegramFake) RoundTrip(request *http.Request) (*http.Response, error) {
 105         response := httptest.NewRecorder()
 106         f.serve(response, request)
 107
 108         ctx := request.Context()
 109         if ctx.Err() != nil {
 110                 return nil, ctx.Err()
 111         }
 112
 113         return response.Result(), nil
 114 }
 115
 116 func (f *TelegramFake) serve(w http.ResponseWriter, r *http.Request) {
 117         id, file := strings.CutPrefix(r.URL.Path, "/file/bot"+telegramTestToken+"/")
 118         if file {
 119                 f.mutex.Lock()
 120
 121                 data, exists := f.images[id]
 122                 f.mutex.Unlock()
 123
 124                 if !exists {
 125                         http.NotFound(w, r)
 126
 127                         return
 128                 }
 129
 130                 _, _ = w.Write(data)
 131
 132                 return
 133         }
 134
 135         method, found := strings.CutPrefix(r.URL.Path, "/bot"+telegramTestToken+"/")
 136         if !found {
 137                 http.Error(w, `{"ok":false,"error_code":404,"description":"Not Found"}`, http.StatusNotFound)
 138
 139                 return
 140         }
 141
 142         var params map[string]any
 143
 144         var err error
 145
 146         if strings.HasPrefix(r.Header.Get("Content-Type"), "multipart/form-data;") {
 147                 err = r.ParseMultipartForm(imageMaxBytes + (1 << 20))
 148                 if err == nil {
 149                         params = make(map[string]any)
 150                         for key, values := range r.MultipartForm.Value {
 151                                 params[key] = values[0]
 152                         }
 153
 154                         for key, files := range r.MultipartForm.File {
 155                                 var file multipart.File
 156
 157                                 file, err = files[0].Open()
 158                                 if err != nil {
 159                                         break
 160                                 }
 161
 162                                 params[key], err = io.ReadAll(file)
 163                                 err = errors.Join(err, file.Close())
 164                                 params[key+"_name"] = files[0].Filename
 165                                 if err != nil {
 166                                         break
 167                                 }
 168                         }
 169
 170                         err = errors.Join(err, r.MultipartForm.RemoveAll())
 171                 }
 172         } else {
 173                 err = json.NewDecoder(r.Body).Decode(&params)
 174         }
 175
 176         if err != nil {
 177                 http.Error(w, err.Error(), http.StatusBadRequest)
 178
 179                 return
 180         }
 181
 182         call := TelegramFakeCall{
 183                 method: method,
 184                 params: params,
 185         }
 186
 187         f.mutex.Lock()
 188
 189         status, description := 0, ""
 190         if f.fail != nil {
 191                 status, description = f.fail(call)
 192         }
 193
 194         stall := f.stall != nil && f.stall(call)
 195         f.mutex.Unlock()
 196
 197         if stall {
 198                 f.record(call)
 199
 200                 select {
 201                 case <-r.Context().Done():
 202                 case <-f.stopped:
 203                 }
 204
 205                 return
 206         }
 207
 208         if status != 0 {
 209                 f.record(call)
 210                 w.WriteHeader(status)
 211
 212                 wait := 0
 213                 if status == http.StatusTooManyRequests {
 214                         wait = 1
 215                 }
 216
 217                 _, _ = fmt.Fprintf(w, `{"ok":false,"error_code":%d,"description":%q,"parameters":{"retry_after":%d}}`, status, description, wait)
 218
 219                 return
 220         }
 221
 222         call.ok = true
 223         f.record(call)
 224
 225         if method == "getFile" {
 226                 _ = json.NewEncoder(w).Encode(map[string]any{
 227                         "ok":     true,
 228                         "result": map[string]any{"file_path": params["file_id"]},
 229                 })
 230
 231                 return
 232         }
 233
 234         if method != "getUpdates" {
 235                 _, _ = io.WriteString(w, `{"ok":true,"result":true}`)
 236                 return
 237         }
 238
 239         offset := int64(params["offset"].(float64))
 240         updates := f.from(offset)
 241         if len(updates) == 0 {
 242                 select {
 243                 case <-f.added:
 244                 case <-time.After(50 * time.Millisecond):
 245                 case <-r.Context().Done():
 246                 }
 247
 248                 updates = f.from(offset)
 249         }
 250
 251         _ = json.NewEncoder(w).Encode(map[string]any{
 252                 "ok":     true,
 253                 "result": updates,
 254         })
 255 }
 256
 257 // record keeps call.
 258 func (f *TelegramFake) record(call TelegramFakeCall) {
 259         f.mutex.Lock()
 260         defer f.mutex.Unlock()
 261
 262         f.calls = append(f.calls, call)
 263 }
 264
 265 // from returns the updates from offset on, and forgets the ones before it, as
 266 // Telegram does.
 267 func (f *TelegramFake) from(offset int64) []TelegramUpdate {
 268         f.mutex.Lock()
 269         defer f.mutex.Unlock()
 270
 271         var updates []TelegramUpdate
 272
 273         kept := f.updates[:0]
 274         for _, update := range f.updates {
 275                 if update.UpdateID >= offset {
 276                         updates = append(updates, update)
 277                         kept = append(kept, update)
 278                 }
 279         }
 280
 281         f.updates = kept
 282
 283         return updates
 284 }
 285
 286 // newTelegramFake starts a fake Bot API server that stops at the end of the
 287 // test.
 288 func newTelegramFake(t *testing.T) *TelegramFake {
 289         t.Helper()
 290
 291         f := &TelegramFake{
 292                 added:   make(chan struct{}, 1),
 293                 stopped: make(chan struct{}),
 294         }
 295
 296         f.Server = httptest.NewServer(http.HandlerFunc(f.serve))
 297         t.Cleanup(f.Close)
 298         t.Cleanup(func() {
 299                 close(f.stopped)
 300         })
 301
 302         return f
 303 }
 304
 305 // newTelegramLocalFake returns the same fake API without network goroutines.
 306 // Its RoundTrip handles requests directly for tests with fake time.
 307 func newTelegramLocalFake(t *testing.T) *TelegramFake {
 308         t.Helper()
 309
 310         f := &TelegramFake{
 311                 Server:  &httptest.Server{URL: "https://example.com"},
 312                 added:   make(chan struct{}, 1),
 313                 stopped: make(chan struct{}),
 314         }
 315         t.Cleanup(func() {
 316                 close(f.stopped)
 317         })
 318
 319         return f
 320 }
 321
 322 // telegramOwnerMessage returns a text message of the owner in the private chat.
 323 func telegramOwnerMessage(text string) *TelegramMessage {
 324         return &TelegramMessage{
 325                 Chat: TelegramChat{
 326                         ID:   telegramTestOwner,
 327                         Type: "private",
 328                 },
 329                 From: &TelegramUser{
 330                         ID: telegramTestOwner,
 331                 },
 332                 Text: text,
 333         }
 334 }