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) Start(t *testing.T, indices ...int) {
  62         t.Helper()
  63
  64         ctx, cancel := context.WithCancel(t.Context())
  65         r.cancel = cancel
  66         for _, index := range indices {
  67                 done := make(chan error, 1)
  68                 r.done = append(r.done, done)
  69                 go func() {
  70                         done <- r.loops[index].Run(ctx)
  71                 }()
  72         }
  73 }
  74
  75 func (r *DaemonMessageTest) Queue(t *testing.T) DaemonMessageFile {
  76         t.Helper()
  77
  78         data, err := os.ReadFile(r.messages.path)
  79         require.NoError(t, err)
  80
  81         var file DaemonMessageFile
  82
  83         require.NoError(t, json.Unmarshal(data, &file))
  84
  85         return file
  86 }
  87
  88 func (r *DaemonMessageTest) Close(t *testing.T) {
  89         t.Helper()
  90
  91         if r.cancel != nil {
  92                 r.cancel()
  93
  94                 for _, done := range r.done {
  95                         require.ErrorIs(t, <-done, context.Canceled)
  96                 }
  97
  98                 r.cancel = nil
  99                 r.done = nil
 100         }
 101
 102         require.NoError(t, r.agents.Close())
 103
 104         r.agents = nil
 105 }