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
  30                 for _, char := range message.Text {
  31                         require.True(t, char == '\n' || (char >= ' ' && char <= '~'), "character %q", char)
  32                 }
  33
  34                 if len(message.Text) >= 30000 {
  35                         long++
  36                 }
  37
  38                 if message.Kind == "work" {
  39                         require.True(t, strings.HasPrefix(message.Text, "["), message.Text)
  40                 }
  41         }
  42
  43         require.Greater(t, counts["user"], 20)
  44         require.Greater(t, counts["tool"], 300)
  45         require.Greater(t, counts["echo"], 300)
  46         require.Greater(t, long, 2)
  47         require.Equal(t, messages, probeLog(probeMessages))
  48
  49         path := os.Getenv("ONECHAT_PROBE_LOG")
  50         if path == "" {
  51                 return
  52         }
  53
  54         var data strings.Builder
  55
  56         for _, message := range messages {
  57                 line, err := json.Marshal(message)
  58                 require.NoError(t, err)
  59
  60                 data.Write(line)
  61                 data.WriteByte('\n')
  62         }
  63
  64         require.NoError(t, os.WriteFile(path, []byte(data.String()), 0o600))
  65 }
  66
  67 func TestProbeReadIgnoresSourceDates(t *testing.T) {
  68         path := filepath.Join(t.TempDir(), "probe.jsonl")
  69         require.NoError(t, os.WriteFile(path, []byte(`{"kind":"user","text":"A probe message","date":"2025-01-01T00:00:00Z"}`+"\n"), 0o600))
  70
  71         entries, err := probeReadLog(path, 1)
  72         require.NoError(t, err)
  73         require.Len(t, entries, 1)
  74         require.Empty(t, entries[0].Date)
  75 }
  76
  77 func TestSlotSummarizerLimitsTheBuildsAtOnce(t *testing.T) {
  78         inner := newFakeSummarizer()
  79         inner.gate = make(chan struct{})
  80         summarizer := ProbeSlotSummarizer{
  81                 inner: inner,
  82                 slots: make(chan struct{}, 2),
  83         }
  84
  85         var group sync.WaitGroup
  86
  87         for n := range 6 {
  88                 group.Go(func() {
  89                         _, err := summarizer.Summarize(t.Context(), Task{
  90                                 Level: 0,
  91                                 Index: n,
  92                                 Input: []string{"user: " + strings.Repeat("x", 600)},
  93                         })
  94                         require.NoError(t, err)
  95                 })
  96         }
  97
  98         require.Eventually(t, func() bool {
  99                 return inner.Active() == 2
 100         }, 10*time.Second, time.Millisecond)
 101
 102         close(inner.gate)
 103         group.Wait()
 104         require.Equal(t, 2, inner.peak)
 105 }