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