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

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