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

   1 package main
   2
   3 import (
   4         "encoding/json"
   5         "os"
   6         "path/filepath"
   7         "strings"
   8         "sync"
   9         "testing"
  10         "time"
  11
  12         "github.com/stretchr/testify/require"
  13 )
  14
  15 // TestProbeLogHasTheSyntheticMix checks that the generated log has readable
  16 // ASCII text with the mix of kinds and sizes of the synthetic log. When
  17 // ONECHAT_PROBE_LOG names a file, the test also writes the log there, one
  18 // message {kind, text} per line, for "onechat probe".
  19 func TestProbeLogHasTheSyntheticMix(t *testing.T) {
  20         messages := probeLog(probeMessages)
  21         require.Len(t, messages, probeMessages)
  22
  23         counts := make(map[string]int)
  24         long := 0
  25         for _, message := range messages {
  26                 counts[message.Kind]++
  27                 require.True(t, logKinds[message.Kind], message.Kind)
  28                 require.NotEmpty(t, message.Text)
  29                 for _, char := range message.Text {
  30                         require.True(t, char == '\n' || (char >= ' ' && char <= '~'), "character %q", char)
  31                 }
  32
  33                 if len(message.Text) >= 30000 {
  34                         long++
  35                 }
  36
  37                 if message.Kind == "work" {
  38                         require.True(t, strings.HasPrefix(message.Text, "["), message.Text)
  39                 }
  40         }
  41
  42         require.Greater(t, counts["user"], 20)
  43         require.Greater(t, counts["tool"], 300)
  44         require.Greater(t, counts["echo"], 300)
  45         require.Greater(t, long, 2)
  46         require.Equal(t, messages, probeLog(probeMessages))
  47
  48         path := os.Getenv("ONECHAT_PROBE_LOG")
  49         if path == "" {
  50                 return
  51         }
  52
  53         var data strings.Builder
  54         for _, message := range messages {
  55                 line, err := json.Marshal(message)
  56                 require.NoError(t, err)
  57
  58                 data.Write(line)
  59                 data.WriteByte('\n')
  60         }
  61
  62         require.NoError(t, os.WriteFile(path, []byte(data.String()), 0o600))
  63 }
  64
  65 func TestProbeReadIgnoresSourceDates(t *testing.T) {
  66         path := filepath.Join(t.TempDir(), "probe.jsonl")
  67         require.NoError(t, os.WriteFile(path, []byte(`{"kind":"user","text":"A probe message","date":"2025-01-01T00:00:00Z"}`+"\n"), 0o600))
  68
  69         entries, err := probeReadLog(path, 1)
  70         require.NoError(t, err)
  71         require.Len(t, entries, 1)
  72         require.Empty(t, entries[0].Date)
  73 }
  74
  75 func TestSlotSummarizerLimitsTheBuildsAtOnce(t *testing.T) {
  76         inner := newFakeSummarizer()
  77         inner.gate = make(chan struct{})
  78         summarizer := ProbeSlotSummarizer{
  79                 inner: inner,
  80                 slots: make(chan struct{}, 2),
  81         }
  82
  83         var group sync.WaitGroup
  84         for n := range 6 {
  85                 group.Go(func() {
  86                         _, err := summarizer.Summarize(t.Context(), Task{
  87                                 Level: 0,
  88                                 Index: n,
  89                                 Input: []string{"user: " + strings.Repeat("x", 600)},
  90                         })
  91                         require.NoError(t, err)
  92                 })
  93         }
  94
  95         require.Eventually(t, func() bool {
  96                 return inner.Active() == 2
  97         }, 10*time.Second, time.Millisecond)
  98
  99         close(inner.gate)
 100         group.Wait()
 101         require.Equal(t, 2, inner.peak)
 102 }