]> git.giacomolli.uk - onechat/blob - telegram_test.go
Name and order the ChatGPT login code
[onechat] / telegram_test.go
1 package main
2
3 import (
4         "log"
5         "net/http"
6         "os"
7         "path/filepath"
8         "slices"
9         "strings"
10         "testing"
11         "time"
12
13         "github.com/stretchr/testify/require"
14 )
15
16 func TestOwnerTalksToTheAgentThroughTelegram(t *testing.T) {
17         server := newModelFake(t)
18         server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
19                 modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("Hello from the agent.")}, 100, 0, 5))
20         }
21
22         telegram := newTelegramFake(t)
23         bot, _ := telegramStartTestBot(t, telegram)
24         socket := startTestDaemonWithBot(t, server, bot)
25         terminal := startTestClient(t, socket)
26         terminal.WaitFor(t, "[connected to OneChat]")
27
28         telegram.Add(TelegramUpdate{
29                 UpdateID: 1,
30                 Message:  telegramOwnerMessage("hi from the phone"),
31         })
32
33         terminal.WaitFor(t, "[telegram] hi from the phone\n[working]\nHello from the agent.\n[done]")
34         telegramWaitForSent(t, telegram, "Hello from the agent.")
35         require.True(t, slices.ContainsFunc(telegram.Calls(), func(call TelegramFakeCall) bool {
36                 return call.method == "sendChatAction" && call.params["action"] == "typing"
37         }))
38
39         terminal.Type("hi from the desk")
40         require.Equal(t, []string{"Hello from the agent.", "[terminal] hi from the desk", "Hello from the agent."}, telegramWaitForSent(t, telegram, "Hello from the agent.", "[terminal] hi from the desk", "Hello from the agent."))
41 }
42
43 func TestBotIgnoresEveryoneButTheOwner(t *testing.T) {
44         telegram := newTelegramFake(t)
45         bot, logs := telegramStartTestBot(t, telegram)
46         socket := startTestDaemonWithBot(t, newModelFake(t), bot)
47         terminal := startTestClient(t, socket)
48         terminal.WaitFor(t, "[connected to OneChat]")
49
50         group := telegramOwnerMessage("secret in a group")
51         group.Chat = TelegramChat{
52                 ID:   -100,
53                 Type: "group",
54         }
55
56         stranger := telegramOwnerMessage("secret of a stranger")
57         stranger.Chat.ID = 77
58         stranger.From.ID = 77
59
60         unsigned := telegramOwnerMessage("secret with no sender")
61         unsigned.From = nil
62
63         for id, message := range []*TelegramMessage{group, stranger, unsigned, nil} {
64                 telegram.Add(TelegramUpdate{
65                         UpdateID: int64(10 + id),
66                         Message:  message,
67                 })
68         }
69
70         require.Eventually(t, func() bool {
71                 offset, err := telegramLoadOffset(bot.offsetPath)
72                 return err == nil && offset == 14
73         }, 10*time.Second, time.Millisecond)
74
75         require.Empty(t, telegram.Sent())
76         require.NotContains(t, terminal.Output(), "secret")
77         require.NotContains(t, logs.String(), "secret")
78         require.Contains(t, logs.String(), "update ignored")
79 }
80
81 func TestBotAnswersThatOnlyTextIsSupported(t *testing.T) {
82         telegram := newTelegramFake(t)
83         bot, _ := telegramStartTestBot(t, telegram)
84         startTestDaemonWithBot(t, newModelFake(t), bot)
85
86         telegram.Add(TelegramUpdate{
87                 UpdateID: 1,
88                 Message:  telegramOwnerMessage(""),
89         })
90
91         telegramWaitForSent(t, telegram, telegramTextOnly)
92 }
93
94 func TestStopFromTelegramInterruptsTheTurn(t *testing.T) {
95         server := newModelFake(t)
96         arrived := make(chan struct{})
97         release := make(chan struct{})
98         t.Cleanup(func() {
99                 close(release)
100         })
101
102         server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
103                 close(arrived)
104                 <-release
105         }
106
107         telegram := newTelegramFake(t)
108         bot, _ := telegramStartTestBot(t, telegram)
109         startTestDaemonWithBot(t, server, bot)
110
111         telegram.Add(TelegramUpdate{
112                 UpdateID: 1,
113                 Message:  telegramOwnerMessage("a long task"),
114         })
115
116         <-arrived
117         telegram.Add(TelegramUpdate{
118                 UpdateID: 2,
119                 Message:  telegramOwnerMessage(stopCommand),
120         })
121
122         telegramWaitForSent(t, telegram, "Error: interrupted")
123
124         telegram.Add(TelegramUpdate{
125                 UpdateID: 3,
126                 Message:  telegramOwnerMessage(stopCommand),
127         })
128
129         telegramWaitForSent(t, telegram, "Error: interrupted", "No turn runs.")
130 }
131
132 func TestBotStartsAtTheStoredOffset(t *testing.T) {
133         server := newModelFake(t)
134         server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
135                 modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("ok")}, 100, 0, 5))
136         }
137
138         telegram := newTelegramFake(t)
139         bot, _ := telegramStartTestBot(t, telegram)
140         require.NoError(t, telegramSaveOffset(bot.offsetPath, 10))
141
142         telegram.Add(TelegramUpdate{
143                 UpdateID: 9,
144                 Message:  telegramOwnerMessage("handled before the restart"),
145         })
146
147         telegram.Add(TelegramUpdate{
148                 UpdateID: 10,
149                 Message:  telegramOwnerMessage("new after the restart"),
150         })
151
152         socket := startTestDaemonWithBot(t, server, bot)
153         terminal := startTestClient(t, socket)
154         terminal.WaitFor(t, "[connected to OneChat]")
155         telegramWaitForSent(t, telegram, "ok")
156
157         require.Equal(t, float64(10), telegram.Calls()[0].params["offset"])
158         require.NotContains(t, terminal.Output(), "handled before the restart")
159
160         require.Eventually(t, func() bool {
161                 offset, err := telegramLoadOffset(bot.offsetPath)
162                 return err == nil && offset == 11
163         }, 10*time.Second, time.Millisecond)
164 }
165
166 func TestBotTriesAgainAfterFailures(t *testing.T) {
167         server := newModelFake(t)
168         server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
169                 modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("ok")}, 100, 0, 5))
170         }
171
172         telegram := newTelegramFake(t)
173         failures := map[string]int{
174                 "getUpdates":      2,
175                 "sendRichMessage": 1,
176         }
177
178         telegram.fail = func(call TelegramFakeCall) (int, string) {
179                 if failures[call.method] == 0 {
180                         return 0, ""
181                 }
182
183                 failures[call.method]--
184                 if call.method == "sendRichMessage" {
185                         return http.StatusTooManyRequests, "Too Many Requests: retry after 1"
186                 }
187
188                 return http.StatusBadGateway, "Bad Gateway"
189         }
190
191         bot, logs := telegramStartTestBot(t, telegram)
192         startTestDaemonWithBot(t, server, bot)
193
194         telegram.Add(TelegramUpdate{
195                 UpdateID: 1,
196                 Message:  telegramOwnerMessage("hello"),
197         })
198
199         telegramWaitForSent(t, telegram, "ok")
200         require.Contains(t, logs.String(), "answered 502")
201         require.Contains(t, logs.String(), "answered 429")
202         require.NotContains(t, logs.String(), telegramTestToken)
203 }
204
205 func TestDeliverTriesAgainAfterATimeout(t *testing.T) {
206         telegram := newTelegramFake(t)
207         stalls := 1
208         telegram.stall = func(call TelegramFakeCall) bool {
209                 if call.method != "sendMessage" || stalls == 0 {
210                         return false
211                 }
212
213                 stalls--
214
215                 return true
216         }
217
218         bot, logs := telegramStartTestBot(t, telegram)
219         bot.api.sendTimeout = 50 * time.Millisecond
220         bot.deliver(t.Context(), TelegramPart{
221                 text: "hello",
222         })
223
224         require.Equal(t, []string{"hello"}, telegram.Sent())
225         require.Len(t, telegram.Calls(), 2)
226         require.Contains(t, logs.String(), "no valid answer")
227         require.NotContains(t, logs.String(), telegramTestToken)
228 }
229
230 func TestRepliesGoAsRichMessages(t *testing.T) {
231         reply := "**Done.** The table:\n\n| Name | Size |\n|:-----|-----:|\n| `a.go` | 10 |\n\n- one\n- two\n\n```go\nx := 1\n```"
232         server := newModelFake(t)
233         server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
234                 modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem(reply)}, 100, 0, 5))
235         }
236
237         telegram := newTelegramFake(t)
238         bot, _ := telegramStartTestBot(t, telegram)
239         socket := startTestDaemonWithBot(t, server, bot)
240         terminal := startTestClient(t, socket)
241         terminal.WaitFor(t, "[connected to OneChat]")
242
243         terminal.Type("make **it** so")
244         telegramWaitForSent(t, telegram, "[terminal] make **it** so", reply)
245
246         telegram.Add(TelegramUpdate{
247                 UpdateID: 1,
248                 Message:  telegramOwnerMessage(""),
249         })
250
251         telegramWaitForSent(t, telegram, telegramTextOnly)
252
253         var sends []TelegramFakeCall
254         for _, call := range telegram.Calls() {
255                 if call.method == "sendMessage" || call.method == "sendRichMessage" {
256                         sends = append(sends, call)
257                 }
258         }
259
260         require.Equal(t, []TelegramFakeCall{
261                 {
262                         method: "sendMessage",
263                         params: map[string]any{
264                                 "chat_id": float64(telegramTestOwner),
265                                 "text":    "[terminal] make **it** so",
266                         },
267                         ok: true,
268                 },
269                 {
270                         method: "sendRichMessage",
271                         params: map[string]any{
272                                 "chat_id": float64(telegramTestOwner),
273                                 "rich_message": map[string]any{
274                                         "markdown": reply,
275                                 },
276                         },
277                         ok: true,
278                 },
279                 {
280                         method: "sendMessage",
281                         params: map[string]any{
282                                 "chat_id": float64(telegramTestOwner),
283                                 "text":    telegramTextOnly,
284                         },
285                         ok: true,
286                 },
287         }, sends)
288 }
289
290 func TestRichMessageFailures(t *testing.T) {
291         testCases := []struct {
292                 name    string
293                 status  int
294                 always  bool
295                 methods []string
296                 sent    []string
297                 logText string
298         }{
299                 {
300                         name:    "bad request sends the reply once with entities",
301                         status:  http.StatusBadRequest,
302                         always:  true,
303                         methods: []string{"sendRichMessage", "sendMessage"},
304                         sent:    []string{"a secret word"},
305                         logText: "rich message refused",
306                 },
307                 {
308                         name:    "unknown method sends the reply once with entities",
309                         status:  http.StatusNotFound,
310                         always:  true,
311                         methods: []string{"sendRichMessage", "sendMessage"},
312                         sent:    []string{"a secret word"},
313                         logText: "rich message refused",
314                 },
315                 {
316                         name:    "rate limit tries the rich message again",
317                         status:  http.StatusTooManyRequests,
318                         methods: []string{"sendRichMessage", "sendRichMessage"},
319                         sent:    []string{"a **secret** word"},
320                         logText: "answered 429",
321                 },
322                 {
323                         name:    "server error tries the rich message again",
324                         status:  http.StatusInternalServerError,
325                         methods: []string{"sendRichMessage", "sendRichMessage"},
326                         sent:    []string{"a **secret** word"},
327                         logText: "answered 500",
328                 },
329                 {
330                         name:    "other refusal drops the reply",
331                         status:  http.StatusForbidden,
332                         always:  true,
333                         methods: []string{"sendRichMessage"},
334                         logText: "message dropped",
335                 },
336         }
337
338         for _, tc := range testCases {
339                 t.Run(tc.name, func(t *testing.T) {
340                         telegram := newTelegramFake(t)
341                         failed := false
342                         telegram.fail = func(call TelegramFakeCall) (int, string) {
343                                 if call.method != "sendRichMessage" || failed && !tc.always {
344                                         return 0, ""
345                                 }
346
347                                 failed = true
348
349                                 return tc.status, "refused for bot" + telegramTestToken
350                         }
351
352                         bot, logs := telegramStartTestBot(t, telegram)
353                         bot.post(t.Context(), Event{
354                                 Kind: "reply",
355                                 Text: "a **secret** word",
356                         })
357
358                         var methods []string
359                         for _, call := range telegram.Calls() {
360                                 methods = append(methods, call.method)
361                         }
362
363                         require.Equal(t, tc.methods, methods)
364                         require.Equal(t, tc.sent, telegram.Sent())
365                         require.Equal(t, 1, strings.Count(logs.String(), "\n"))
366                         require.Contains(t, logs.String(), tc.logText)
367                         require.NotContains(t, logs.String(), telegramTestToken)
368                         require.NotContains(t, logs.String(), "secret word")
369
370                         calls := telegram.Calls()
371                         if calls[len(calls)-1].method == "sendMessage" {
372                                 require.NotNil(t, calls[len(calls)-1].params["entities"])
373                         }
374                 })
375         }
376 }
377
378 func TestLongReplyGoesAsRichMessagesSplitBetweenBlocks(t *testing.T) {
379         paragraphs := strings.Repeat("A paragraph of the reply.\n\n", telegramRichBytes/len("A paragraph of the reply.\n\n"))
380         table := "| a | b |\n|---|---|\n| 1 | 2 |\n"
381         reply := paragraphs + table
382
383         telegram := newTelegramFake(t)
384         bot, _ := telegramStartTestBot(t, telegram)
385         bot.post(t.Context(), Event{
386                 Kind: "reply",
387                 Text: reply,
388         })
389
390         require.Equal(t, []string{paragraphs, table}, telegram.Sent())
391 }
392
393 func TestRefusedEntitiesGoAgainAsPlainText(t *testing.T) {
394         telegram := newTelegramFake(t)
395         telegram.fail = func(call TelegramFakeCall) (int, string) {
396                 if call.params["entities"] == nil {
397                         return 0, ""
398                 }
399
400                 return http.StatusBadRequest, "Bad Request: can't parse entities of bot" + telegramTestToken
401         }
402
403         bot, logs := telegramStartTestBot(t, telegram)
404         text, entities := telegramMarkdownEntities("a **secret** word")
405         bot.deliver(t.Context(), TelegramPart{
406                 text:     text,
407                 entities: entities,
408         })
409
410         calls := telegram.Calls()
411         require.Len(t, calls, 2)
412         require.NotNil(t, calls[0].params["entities"])
413         require.Nil(t, calls[1].params["entities"])
414         require.Equal(t, []string{"a secret word"}, telegram.Sent())
415         require.Equal(t, 1, strings.Count(logs.String(), "\n"))
416         require.Contains(t, logs.String(), "entities refused")
417         require.NotContains(t, logs.String(), telegramTestToken)
418         require.NotContains(t, logs.String(), "secret word")
419 }
420
421 func TestBotLogsNoTokenWhenTheNetworkFails(t *testing.T) {
422         telegram := newTelegramFake(t)
423         telegram.Close()
424
425         bot, logs := telegramStartTestBot(t, telegram)
426         startTestDaemonWithBot(t, newModelFake(t), bot)
427
428         require.Eventually(t, func() bool {
429                 return strings.Count(logs.String(), "no valid answer") >= 2
430         }, 10*time.Second, time.Millisecond)
431         require.NotContains(t, logs.String(), telegramTestToken)
432 }
433
434 func TestSplitMessage(t *testing.T) {
435         testCases := []struct {
436                 name  string
437                 text  string
438                 limit int
439                 parts []string
440         }{
441                 {
442                         name:  "short text stays whole",
443                         text:  "one two",
444                         limit: 10,
445                         parts: []string{"one two"},
446                 },
447                 {
448                         name:  "cut after the last line break",
449                         text:  "first line\nsecond line",
450                         limit: 16,
451                         parts: []string{"first line\n", "second line"},
452                 },
453                 {
454                         name:  "cut after the last space without a line break",
455                         text:  "aaaa bbbb cccc",
456                         limit: 12,
457                         parts: []string{"aaaa bbbb ", "cccc"},
458                 },
459                 {
460                         name:  "hard cut of one long word",
461                         text:  "abcdefghij",
462                         limit: 4,
463                         parts: []string{"abcd", "efgh", "ij"},
464                 },
465                 {
466                         name:  "a character outside the BMP counts two units and is not cut",
467                         text:  "ab\U0001F600cd",
468                         limit: 3,
469                         parts: []string{"ab", "\U0001F600c", "d"},
470                 },
471         }
472
473         for _, tc := range testCases {
474                 t.Run(tc.name, func(t *testing.T) {
475                         require.Equal(t, tc.parts, telegramSplitMessage(tc.text, tc.limit))
476                 })
477         }
478
479         long := strings.Repeat("word ", 2000)
480         parts := telegramSplitMessage(long, telegramMessageUnits)
481         require.Len(t, parts, 3)
482         require.Equal(t, long, strings.Join(parts, ""))
483 }
484
485 func TestSplitEntitiesCutsAnEntityAcrossASplit(t *testing.T) {
486         parts := telegramSplitEntities("aabbbbcc", []TelegramEntity{
487                 {
488                         Type:   "bold",
489                         Offset: 2,
490                         Length: 4,
491                 },
492                 {
493                         Type:   "italic",
494                         Offset: 6,
495                         Length: 2,
496                 },
497         }, 5)
498
499         require.Equal(t, []TelegramPart{
500                 {
501                         text: "aabbb",
502                         entities: []TelegramEntity{
503                                 {
504                                         Type:   "bold",
505                                         Offset: 2,
506                                         Length: 3,
507                                 },
508                         },
509                 },
510                 {
511                         text: "bcc",
512                         entities: []TelegramEntity{
513                                 {
514                                         Type:   "bold",
515                                         Offset: 0,
516                                         Length: 1,
517                                 },
518                                 {
519                                         Type:   "italic",
520                                         Offset: 1,
521                                         Length: 2,
522                                 },
523                         },
524                 },
525         }, parts)
526 }
527
528 func TestSplitRich(t *testing.T) {
529         testCases := []struct {
530                 name     string
531                 markdown string
532                 limit    int
533                 parts    []string
534         }{
535                 {
536                         name:     "short reply stays whole",
537                         markdown: "| a | b |",
538                         limit:    100,
539                         parts:    []string{"| a | b |"},
540                 },
541                 {
542                         name:     "cut after a blank line keeps a table whole",
543                         markdown: "Intro.\n\n| a | b |\n|---|---|\n| 1 | 2 |\n",
544                         limit:    30,
545                         parts:    []string{"Intro.\n\n", "| a | b |\n|---|---|\n| 1 | 2 |\n"},
546                 },
547                 {
548                         name:     "blank line in a code block is not a cut",
549                         markdown: "Text.\n\n```go\nx\n\ny\n```\n\nEnd.\n",
550                         limit:    20,
551                         parts:    []string{"Text.\n\n", "```go\nx\n\ny\n```\n\n", "End.\n"},
552                 },
553                 {
554                         name:     "paragraph with no blank line is cut after a line",
555                         markdown: "aaa\nbbb\nccc\n",
556                         limit:    9,
557                         parts:    []string{"aaa\nbbb\n", "ccc\n"},
558                 },
559                 {
560                         name:     "long code block is closed and opened again",
561                         markdown: "```sh\none\ntwo\nthree\n```\n",
562                         limit:    16,
563                         parts:    []string{"```sh\none\n```\n", "```sh\ntwo\n```\n", "```sh\nthree\n```\n"},
564                 },
565                 {
566                         name:     "long line is cut at the limit and not inside a character",
567                         markdown: "ab\u00e9cd",
568                         limit:    3,
569                         parts:    []string{"ab", "\u00e9c", "d"},
570                 },
571         }
572
573         for _, tc := range testCases {
574                 t.Run(tc.name, func(t *testing.T) {
575                         require.Equal(t, tc.parts, telegramSplitRich(tc.markdown, tc.limit))
576                 })
577         }
578
579         long := strings.Repeat("A paragraph of the reply.\n\n", 2000)
580         parts := telegramSplitRich(long, telegramRichBytes)
581         require.Len(t, parts, 2)
582         require.Len(t, parts[0], telegramRichBytes-telegramRichBytes%len("A paragraph of the reply.\n\n"))
583         require.Equal(t, long, strings.Join(parts, ""))
584 }
585
586 func TestLoadBotConfig(t *testing.T) {
587         testCases := []struct {
588                 name      string
589                 content   string
590                 mode      os.FileMode
591                 errorText string
592         }{
593                 {
594                         name:    "owner file is read",
595                         content: `{"token": "` + telegramTestToken + `", "owner": 4242}`,
596                         mode:    0o600,
597                 },
598                 {
599                         name:      "file that others can read is refused",
600                         content:   `{"token": "` + telegramTestToken + `", "owner": 4242}`,
601                         mode:      0o640,
602                         errorText: "mode must be 0600",
603                 },
604                 {
605                         name:      "token with a path character is refused",
606                         content:   `{"token": "123:abc/../x", "owner": 4242}`,
607                         mode:      0o600,
608                         errorText: "form of a bot token",
609                 },
610                 {
611                         name:      "missing owner is refused",
612                         content:   `{"token": "` + telegramTestToken + `"}`,
613                         mode:      0o600,
614                         errorText: "numeric Telegram user ID",
615                 },
616                 {
617                         name:      "unknown field is refused",
618                         content:   `{"token": "` + telegramTestToken + `", "owner": 4242, "chat": 1}`,
619                         mode:      0o600,
620                         errorText: "not valid JSON",
621                 },
622         }
623
624         for _, tc := range testCases {
625                 t.Run(tc.name, func(t *testing.T) {
626                         dir := t.TempDir()
627                         path := filepath.Join(dir, telegramConfigFileName)
628                         require.NoError(t, os.WriteFile(path, []byte(tc.content), tc.mode))
629                         require.NoError(t, os.Chmod(path, tc.mode))
630
631                         config, found, err := loadTelegramConfig(dir)
632                         if tc.errorText != "" {
633                                 require.ErrorContains(t, err, tc.errorText)
634                                 require.NotContains(t, err.Error(), telegramTestToken)
635
636                                 return
637                         }
638
639                         require.NoError(t, err)
640                         require.True(t, found)
641                         require.Equal(t, TelegramConfig{
642                                 Token: telegramTestToken,
643                                 Owner: telegramTestOwner,
644                         }, config)
645                 })
646         }
647
648         _, found, err := loadTelegramConfig(t.TempDir())
649         require.NoError(t, err)
650         require.False(t, found)
651 }
652
653 // telegramStartTestBot returns a bot of telegramTestOwner on server, with its
654 // offset file in a temporary directory and short waits, and the buffer of its
655 // log.
656 func telegramStartTestBot(t *testing.T, server *TelegramFake) (*TelegramBot, *syncBuffer) {
657         t.Helper()
658
659         logs := &syncBuffer{}
660         bot := newTelegramBot(newTelegramAPI(server.URL, telegramTestToken), telegramTestOwner, filepath.Join(t.TempDir(), telegramOffsetFileName), log.New(logs, "", 0))
661         bot.minWait = time.Millisecond
662         bot.maxWait = 10 * time.Millisecond
663
664         return bot, logs
665 }
666
667 // telegramWaitForSent waits until the bot sent the texts in order, and returns
668 // them.
669 func telegramWaitForSent(t *testing.T, server *TelegramFake, texts ...string) []string {
670         t.Helper()
671
672         require.Eventually(t, func() bool {
673                 sent := server.Sent()
674                 return len(sent) >= len(texts) && slices.Equal(sent[len(sent)-len(texts):], texts)
675         }, 10*time.Second, time.Millisecond, "sent so far: %q", server.Sent())
676
677         return server.Sent()
678 }