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

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