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

   1 package main
   2
   3 import (
   4         "context"
   5         "encoding/json"
   6         "errors"
   7         "os"
   8         "path/filepath"
   9         "testing"
  10
  11         "github.com/stretchr/testify/require"
  12 )
  13
  14 // DaemonMessageTest owns test agents and the loops selected by Start.
  15 type DaemonMessageTest struct {
  16         agents   DaemonAgents
  17         loops    []*TurnLoop
  18         events   []*TurnEventRecorder
  19         messages *DaemonMessages
  20         logs     *SyncBuffer
  21         cancel   context.CancelFunc
  22         done     []chan error
  23 }
  24
  25 func (r *DaemonMessageTest) Start(t *testing.T, indices ...int) {
  26         t.Helper()
  27
  28         ctx, cancel := context.WithCancel(t.Context())
  29         r.cancel = cancel
  30         for _, index := range indices {
  31                 done := make(chan error, 1)
  32                 r.done = append(r.done, done)
  33                 go func() {
  34                         done <- r.loops[index].Run(ctx)
  35                 }()
  36         }
  37 }
  38
  39 func (r *DaemonMessageTest) Queue(t *testing.T) DaemonMessageFile {
  40         t.Helper()
  41
  42         data, err := os.ReadFile(r.messages.path)
  43         require.NoError(t, err)
  44
  45         var file DaemonMessageFile
  46
  47         require.NoError(t, json.Unmarshal(data, &file))
  48
  49         return file
  50 }
  51
  52 func (r *DaemonMessageTest) Close(t *testing.T) {
  53         t.Helper()
  54
  55         if r.cancel != nil {
  56                 r.cancel()
  57
  58                 for _, done := range r.done {
  59                         require.ErrorIs(t, <-done, context.Canceled)
  60                 }
  61
  62                 r.cancel = nil
  63                 r.done = nil
  64         }
  65
  66         require.NoError(t, r.agents.Close())
  67
  68         r.agents = nil
  69 }
  70
  71 func newDaemonMessageTest(t *testing.T, dir string, server *ModelFake) *DaemonMessageTest {
  72         t.Helper()
  73
  74         _, err := os.Stat(filepath.Join(dir, daemonAgentDirName, "children"))
  75         if errors.Is(err, os.ErrNotExist) {
  76                 daemonChildConfig(t, dir, "Garden")
  77                 daemonChildConfig(t, dir, "Library")
  78         } else {
  79                 require.NoError(t, err)
  80         }
  81
  82         test := &DaemonMessageTest{
  83                 logs: &SyncBuffer{},
  84         }
  85         test.agents, err = openDaemonAgents(t.Context(), dir, server.URL, ModelAPIKey(modelTestKey), defaultConfig().Compactor, newTestClock().Now, daemonTestLogger(test.logs))
  86         require.NoError(t, err)
  87         t.Cleanup(func() {
  88                 test.Close(t)
  89         })
  90
  91         for _, agent := range test.agents {
  92                 events := &TurnEventRecorder{}
  93                 loop, err := agent.Loop(nil, events.Add)
  94                 require.NoError(t, err)
  95
  96                 test.events = append(test.events, events)
  97                 test.loops = append(test.loops, loop)
  98         }
  99
 100         test.messages, err = openDaemonMessages(test.loops)
 101         require.NoError(t, err)
 102
 103         return test
 104 }