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 newDaemonMessageTest(t *testing.T, dir string, server *ModelFake) *DaemonMessageTest {
  26         t.Helper()
  27
  28         _, err := os.Stat(filepath.Join(dir, daemonAgentDirName, "children"))
  29         if errors.Is(err, os.ErrNotExist) {
  30                 daemonChildConfig(t, dir, "Garden")
  31                 daemonChildConfig(t, dir, "Library")
  32         } else {
  33                 require.NoError(t, err)
  34         }
  35
  36         test := &DaemonMessageTest{
  37                 logs: &SyncBuffer{},
  38         }
  39
  40         test.agents, err = openDaemonAgents(t.Context(), dir, server.URL, ModelAPIKey(modelTestKey), defaultConfig().Compactor, newTestClock().Now, daemonTestLogger(test.logs))
  41         require.NoError(t, err)
  42         t.Cleanup(func() {
  43                 test.Close(t)
  44         })
  45
  46         for _, agent := range test.agents {
  47                 events := &TurnEventRecorder{}
  48                 loop, err := agent.Loop(nil, events.Add)
  49                 require.NoError(t, err)
  50
  51                 test.events = append(test.events, events)
  52                 test.loops = append(test.loops, loop)
  53         }
  54
  55         test.messages, err = openDaemonMessages(test.loops)
  56         require.NoError(t, err)
  57
  58         return test
  59 }
  60
  61 func (r *DaemonMessageTest) Close(t *testing.T) {
  62         t.Helper()
  63
  64         if r.cancel != nil {
  65                 r.cancel()
  66
  67                 for _, done := range r.done {
  68                         require.ErrorIs(t, <-done, context.Canceled)
  69                 }
  70
  71                 r.cancel = nil
  72                 r.done = nil
  73         }
  74
  75         require.NoError(t, r.agents.Close())
  76
  77         r.agents = nil
  78 }
  79
  80 func (r *DaemonMessageTest) Start(t *testing.T, indices ...int) {
  81         t.Helper()
  82
  83         ctx, cancel := context.WithCancel(t.Context())
  84         r.cancel = cancel
  85         for _, index := range indices {
  86                 done := make(chan error, 1)
  87                 r.done = append(r.done, done)
  88                 go func() {
  89                         done <- r.loops[index].Run(ctx)
  90                 }()
  91         }
  92 }
  93
  94 func (r *DaemonMessageTest) Queue(t *testing.T) DaemonMessageFile {
  95         t.Helper()
  96
  97         data, err := os.ReadFile(r.messages.path)
  98         require.NoError(t, err)
  99
 100         var file DaemonMessageFile
 101
 102         require.NoError(t, json.Unmarshal(data, &file))
 103
 104         return file
 105 }