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

   1 package main
   2
   3 import (
   4         "context"
   5         "path/filepath"
   6         "testing"
   7         "time"
   8
   9         "github.com/stretchr/testify/require"
  10 )
  11
  12 // daemonStartTestWithBot is daemonStartTest with bot, which can be nil.
  13 // WhatsApp is not linked.
  14 func daemonStartTestWithBot(t *testing.T, server *ModelFake, bot *TelegramBot) string {
  15         t.Helper()
  16
  17         whatsapp, _, _ := whatsappStartTest(t, t.TempDir(), newTestClock())
  18         socket, _ := daemonStartTestWithWhatsApp(t, server, bot, whatsapp)
  19
  20         return socket
  21 }
  22
  23 // daemonStartTestWithWhatsApp is daemonStartTestWithBot with whatsapp. It also
  24 // returns the memory of the agent.
  25 func daemonStartTestWithWhatsApp(t *testing.T, server *ModelFake, bot *TelegramBot, whatsapp *WhatsApp) (string, *Memory) {
  26         t.Helper()
  27
  28         clock := newTestClock()
  29
  30         return daemonStartTestWithSchedules(t, server, bot, whatsapp, newTestSchedules(t, clock.Now), clock.Now)
  31 }
  32
  33 // daemonStartTestWithSchedules is daemonStartTestWithWhatsApp with schedules,
  34 // and with clock for the memory and the turns.
  35 func daemonStartTestWithSchedules(t *testing.T, server *ModelFake, bot *TelegramBot, whatsapp *WhatsApp, schedules *Schedules, clock func() time.Time) (string, *Memory) {
  36         t.Helper()
  37
  38         dir := t.TempDir()
  39         agent, err := openAgent(filepath.Join(dir, daemonAgentDirName))
  40         require.NoError(t, err)
  41
  42         agent.config.Workdir = t.TempDir()
  43         memory, err := openMemory(t.Context(), agent.dir, clock, newFakeSummarizer(), discardLog())
  44         require.NoError(t, err)
  45
  46         socket := filepath.Join(dir, socketFileName)
  47         listener, err := daemonListenSocket(socket)
  48         require.NoError(t, err)
  49
  50         client := newModelClient(server.URL, ModelAPIKey(modelTestKey), agent.CacheKey(), filepath.Join(agent.dir, "usage.jsonl"))
  51         ctx, cancel := context.WithCancel(t.Context())
  52         done := make(chan error)
  53         subagents := newTestSubagents(t)
  54         go func() {
  55                 done <- daemonServe(ctx, listener, bot, whatsapp, schedules, func(notify func(Event)) (*TurnLoop, error) {
  56                         return newTurnLoop(agent, memory, subagents, whatsapp, schedules, client, clock, discardLog(), notify)
  57                 })
  58         }()
  59
  60         t.Cleanup(func() {
  61                 cancel()
  62                 require.NoError(t, <-done)
  63                 require.NoError(t, memory.Close())
  64         })
  65
  66         return socket, memory
  67 }