[onechat] / telegram_fake_helper_test.go

package main

import (
        "encoding/json"
        "fmt"
        "io"
        "net/http"
        "net/http/httptest"
        "strings"
        "sync"
        "testing"
        "time"
)

const (
        // telegramTestToken is the bot token of the tests. It has the form of a
        // real token.
        telegramTestToken = "123456:TEST-secret_token"

        // telegramTestOwner is the user ID of the owner of the test bot.
        telegramTestOwner = 4242
)

// TelegramFakeCall is one call of the fake Bot API. OK tells if the call got a
// success answer.
type TelegramFakeCall struct {
        method string
        params map[string]any
        ok     bool
}

// TelegramFake is a fake Bot API server for telegramTestToken. getUpdates
// returns the added updates from the offset on, or waits a short time when
// there are none. The other methods succeed, unless fail gives an error answer
// for a call. An answer with HTTP 429 asks for a wait of one second. When stall
// is true for a call, the call gets no answer until the client gives it up or
// the test ends. Fail and stall run with the mutex held.
type TelegramFake struct {
        *httptest.Server

        mutex   sync.Mutex
        updates []TelegramUpdate
        calls   []TelegramFakeCall
        fail    func(call TelegramFakeCall) (int, string)
        stall   func(call TelegramFakeCall) bool
        added   chan struct{}
        stopped chan struct{}
        images  map[string][]byte
}

// Add makes update available to getUpdates.
func (f *TelegramFake) Add(update TelegramUpdate) {
        f.mutex.Lock()

        f.updates = append(f.updates, update)
        f.mutex.Unlock()

        select {
        case f.added <- struct{}{}:
        default:
        }
}

// Calls returns the calls so far.
func (f *TelegramFake) Calls() []TelegramFakeCall {
        f.mutex.Lock()
        defer f.mutex.Unlock()

        return append([]TelegramFakeCall(nil), f.calls...)
}

// Sent returns the texts of the sendMessage calls and the Markdown of the
// sendRichMessage calls that succeeded, in order.
func (f *TelegramFake) Sent() []string {
        var texts []string

        for _, call := range f.Calls() {
                switch {
                case !call.ok:
                case call.method == "sendMessage":
                        texts = append(texts, call.params["text"].(string))
                case call.method == "sendRichMessage":
                        texts = append(texts, call.params["rich_message"].(map[string]any)["markdown"].(string))
                }
        }

        return texts
}

func (f *TelegramFake) Image(id string, data []byte) {
        f.mutex.Lock()
        defer f.mutex.Unlock()

        if f.images == nil {
                f.images = make(map[string][]byte)
        }

        f.images[id] = data
}

func (f *TelegramFake) serve(w http.ResponseWriter, r *http.Request) {
        id, file := strings.CutPrefix(r.URL.Path, "/file/bot"+telegramTestToken+"/")
        if file {
                f.mutex.Lock()

                data, exists := f.images[id]
                f.mutex.Unlock()

                if !exists {
                        http.NotFound(w, r)

                        return
                }

                _, _ = w.Write(data)

                return
        }

        method, found := strings.CutPrefix(r.URL.Path, "/bot"+telegramTestToken+"/")
        if !found {
                http.Error(w, `{"ok":false,"error_code":404,"description":"Not Found"}`, http.StatusNotFound)

                return
        }

        var params map[string]any

        err := json.NewDecoder(r.Body).Decode(&params)
        if err != nil {
                http.Error(w, err.Error(), http.StatusBadRequest)

                return
        }

        call := TelegramFakeCall{
                method: method,
                params: params,
        }

        f.mutex.Lock()

        status, description := 0, ""
        if f.fail != nil {
                status, description = f.fail(call)
        }

        stall := f.stall != nil && f.stall(call)
        f.mutex.Unlock()

        if stall {
                f.record(call)

                select {
                case <-r.Context().Done():
                case <-f.stopped:
                }

                return
        }

        if status != 0 {
                f.record(call)
                w.WriteHeader(status)

                wait := 0
                if status == http.StatusTooManyRequests {
                        wait = 1
                }

                _, _ = fmt.Fprintf(w, `{"ok":false,"error_code":%d,"description":%q,"parameters":{"retry_after":%d}}`, status, description, wait)

                return
        }

        call.ok = true
        f.record(call)

        if method == "getFile" {
                _ = json.NewEncoder(w).Encode(map[string]any{
                        "ok":     true,
                        "result": map[string]any{"file_path": params["file_id"]},
                })

                return
        }

        if method != "getUpdates" {
                _, _ = io.WriteString(w, `{"ok":true,"result":true}`)
                return
        }

        offset := int64(params["offset"].(float64))
        updates := f.from(offset)
        if len(updates) == 0 {
                select {
                case <-f.added:
                case <-time.After(50 * time.Millisecond):
                case <-r.Context().Done():
                }

                updates = f.from(offset)
        }

        _ = json.NewEncoder(w).Encode(map[string]any{
                "ok":     true,
                "result": updates,
        })
}

// record keeps call.
func (f *TelegramFake) record(call TelegramFakeCall) {
        f.mutex.Lock()
        defer f.mutex.Unlock()

        f.calls = append(f.calls, call)
}

// from returns the updates from offset on, and forgets the ones before it, as
// Telegram does.
func (f *TelegramFake) from(offset int64) []TelegramUpdate {
        f.mutex.Lock()
        defer f.mutex.Unlock()

        var updates []TelegramUpdate

        kept := f.updates[:0]
        for _, update := range f.updates {
                if update.UpdateID >= offset {
                        updates = append(updates, update)
                        kept = append(kept, update)
                }
        }

        f.updates = kept

        return updates
}

// newTelegramFake starts a fake Bot API server that stops at the end of the
// test.
func newTelegramFake(t *testing.T) *TelegramFake {
        t.Helper()

        f := &TelegramFake{
                added:   make(chan struct{}, 1),
                stopped: make(chan struct{}),
        }

        f.Server = httptest.NewServer(http.HandlerFunc(f.serve))
        t.Cleanup(f.Close)
        t.Cleanup(func() {
                close(f.stopped)
        })

        return f
}

// telegramOwnerMessage returns a text message of the owner in the private chat.
func telegramOwnerMessage(text string) *TelegramMessage {
        return &TelegramMessage{
                Chat: TelegramChat{
                        ID:   telegramTestOwner,
                        Type: "private",
                },
                From: &TelegramUser{
                        ID: telegramTestOwner,
                },
                Text: text,
        }
}