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 }
  49
  50 // Add makes update available to getUpdates.
  51 func (f *TelegramFake) Add(update TelegramUpdate) {
  52         f.mutex.Lock()
  53
  54         f.updates = append(f.updates, update)
  55         f.mutex.Unlock()
  56
  57         select {
  58         case f.added <- struct{}{}:
  59         default:
  60         }
  61 }
  62
  63 // Calls returns the calls so far.
  64 func (f *TelegramFake) Calls() []TelegramFakeCall {
  65         f.mutex.Lock()
  66         defer f.mutex.Unlock()
  67
  68         return append([]TelegramFakeCall(nil), f.calls...)
  69 }
  70
  71 // Sent returns the texts of the sendMessage calls and the Markdown of the
  72 // sendRichMessage calls that succeeded, in order.
  73 func (f *TelegramFake) Sent() []string {
  74         var texts []string
  75
  76         for _, call := range f.Calls() {
  77                 switch {
  78                 case !call.ok:
  79                 case call.method == "sendMessage":
  80                         texts = append(texts, call.params["text"].(string))
  81                 case call.method == "sendRichMessage":
  82                         texts = append(texts, call.params["rich_message"].(map[string]any)["markdown"].(string))
  83                 }
  84         }
  85
  86         return texts
  87 }
  88
  89 func (f *TelegramFake) serve(w http.ResponseWriter, r *http.Request) {
  90         method, found := strings.CutPrefix(r.URL.Path, "/bot"+telegramTestToken+"/")
  91         if !found {
  92                 http.Error(w, `{"ok":false,"error_code":404,"description":"Not Found"}`, http.StatusNotFound)
  93
  94                 return
  95         }
  96
  97         var params map[string]any
  98
  99         err := json.NewDecoder(r.Body).Decode(&params)
 100         if err != nil {
 101                 http.Error(w, err.Error(), http.StatusBadRequest)
 102
 103                 return
 104         }
 105
 106         call := TelegramFakeCall{
 107                 method: method,
 108                 params: params,
 109         }
 110
 111         f.mutex.Lock()
 112
 113         status, description := 0, ""
 114         if f.fail != nil {
 115                 status, description = f.fail(call)
 116         }
 117
 118         stall := f.stall != nil && f.stall(call)
 119         f.mutex.Unlock()
 120
 121         if stall {
 122                 f.record(call)
 123
 124                 select {
 125                 case <-r.Context().Done():
 126                 case <-f.stopped:
 127                 }
 128
 129                 return
 130         }
 131
 132         if status != 0 {
 133                 f.record(call)
 134                 w.WriteHeader(status)
 135
 136                 wait := 0
 137                 if status == http.StatusTooManyRequests {
 138                         wait = 1
 139                 }
 140
 141                 _, _ = fmt.Fprintf(w, `{"ok":false,"error_code":%d,"description":%q,"parameters":{"retry_after":%d}}`, status, description, wait)
 142
 143                 return
 144         }
 145
 146         call.ok = true
 147         f.record(call)
 148
 149         if method != "getUpdates" {
 150                 _, _ = io.WriteString(w, `{"ok":true,"result":true}`)
 151                 return
 152         }
 153
 154         offset := int64(params["offset"].(float64))
 155         updates := f.from(offset)
 156         if len(updates) == 0 {
 157                 select {
 158                 case <-f.added:
 159                 case <-time.After(50 * time.Millisecond):
 160                 case <-r.Context().Done():
 161                 }
 162
 163                 updates = f.from(offset)
 164         }
 165
 166         _ = json.NewEncoder(w).Encode(map[string]any{
 167                 "ok":     true,
 168                 "result": updates,
 169         })
 170 }
 171
 172 // record keeps call.
 173 func (f *TelegramFake) record(call TelegramFakeCall) {
 174         f.mutex.Lock()
 175         defer f.mutex.Unlock()
 176
 177         f.calls = append(f.calls, call)
 178 }
 179
 180 // from returns the updates from offset on, and forgets the ones before it, as
 181 // Telegram does.
 182 func (f *TelegramFake) from(offset int64) []TelegramUpdate {
 183         f.mutex.Lock()
 184         defer f.mutex.Unlock()
 185
 186         var updates []TelegramUpdate
 187
 188         kept := f.updates[:0]
 189         for _, update := range f.updates {
 190                 if update.UpdateID >= offset {
 191                         updates = append(updates, update)
 192                         kept = append(kept, update)
 193                 }
 194         }
 195
 196         f.updates = kept
 197
 198         return updates
 199 }
 200
 201 // newTelegramFake starts a fake Bot API server that stops at the end of the
 202 // test.
 203 func newTelegramFake(t *testing.T) *TelegramFake {
 204         t.Helper()
 205
 206         f := &TelegramFake{
 207                 added:   make(chan struct{}, 1),
 208                 stopped: make(chan struct{}),
 209         }
 210
 211         f.Server = httptest.NewServer(http.HandlerFunc(f.serve))
 212         t.Cleanup(f.Close)
 213         t.Cleanup(func() {
 214                 close(f.stopped)
 215         })
 216
 217         return f
 218 }
 219
 220 // telegramOwnerMessage returns a text message of the owner in the private chat.
 221 func telegramOwnerMessage(text string) *TelegramMessage {
 222         return &TelegramMessage{
 223                 Chat: TelegramChat{
 224                         ID:   telegramTestOwner,
 225                         Type: "private",
 226                 },
 227                 From: &TelegramUser{
 228                         ID: telegramTestOwner,
 229                 },
 230                 Text: text,
 231         }
 232 }