X-Git-Url: https://git.giacomolli.uk/onechat/blobdiff_plain/522d226230a2eae0a8a33ae0fbf32f96c6886b26..c9183f90d757006e7221ecd0d4dc0db2a25f87e0:/probe_test.go?ds=sidebyside diff --git a/probe_test.go b/probe_test.go index c7436dc..0651281 100644 --- a/probe_test.go +++ b/probe_test.go @@ -9,7 +9,9 @@ import ( "path/filepath" "strconv" "strings" + "sync" "testing" + "time" "github.com/stretchr/testify/require" ) @@ -122,22 +124,17 @@ func TestProbeRunStopsWhenTheQuotaGrows(t *testing.T) { writeEvents(w, completedEvent([]any{messageItem(strings.Repeat("s", 400))}, 1000, 800, 50)) } - root := t.TempDir() - usagePath := filepath.Join(root, "usage.jsonl") - client := newClient(server.URL, apiKey(testKey), "onechat-"+probeName, usagePath) - memory, err := openMemory(t.Context(), root, newTestClock().Now, newCompactor(client, probeName, defaultConfig().Compactor)) + run := newTestProbeRun(t, server) + appended, stop := run.Run(t.Context(), probeLog(probeMessages)) + sizes, err := modelLines(run.memory) require.NoError(t, err) + require.NoError(t, run.memory.Close()) - appended, stop := runProbe(t.Context(), memory, probeLog(probeMessages), usagePath) - sizes, err := modelLines(memory) - require.NoError(t, err) - require.NoError(t, memory.Close()) - - require.Contains(t, stop, "x-codex-primary-used-percent grew from 10 to") + require.Regexp(t, `^x-codex-primary-used-percent grew from \d+ to \d+$`, stop) require.Less(t, appended, probeMessages) require.NotEmpty(t, sizes) - records, err := readUsageFile(usagePath) + records, err := readUsageFile(run.usagePath) require.NoError(t, err) require.Greater(t, len(records), int(probeQuotaLimit)) @@ -146,7 +143,153 @@ func TestProbeRunStopsWhenTheQuotaGrows(t *testing.T) { require.Contains(t, report.String(), "stopped early: ") require.Contains(t, report.String(), "cached (80.0%)") require.Contains(t, report.String(), "line bytes: min 400, median 400, mean 400, max 400, 0 over 512") - require.Contains(t, report.String(), "quota at the first call: primary-used-percent=10") + require.Contains(t, report.String(), "calls at once: at most ") + require.Contains(t, report.String(), "quota at the first call: primary-used-percent=") +} + +func TestProbeRunStopsAtTheCallLimit(t *testing.T) { + server := newFakeResponses(t) + run := newTestProbeRun(t, server) + run.client.LimitCalls(5) + + entries := make([]probeEntry, 40) + for n := range entries { + entries[n] = probeEntry{ + Kind: "echo", + Text: strings.Repeat("long output ", 100), + } + } + + _, stop := run.Run(t.Context(), entries) + require.NoError(t, run.memory.Close()) + require.Equal(t, "call limit of 5 reached", stop) + require.Equal(t, 5, server.Calls()) +} + +func TestProbeRunWaitsOnlyBeforeUserMessages(t *testing.T) { + gate := make(chan struct{}) + var mutex sync.Mutex + running, most := 0, 0 + server := newFakeResponses(t) + server.reply = func(w http.ResponseWriter, call fakeCall) { + mutex.Lock() + running++ + most = max(most, running) + mutex.Unlock() + + <-gate + + mutex.Lock() + running-- + mutex.Unlock() + + writeEvents(w, completedEvent([]any{messageItem("a line")}, 100, 0, 5)) + } + + run := newTestProbeRun(t, server) + long := strings.Repeat("long output ", 100) + entries := []probeEntry{ + { + Kind: "echo", + Text: long, + }, + { + Kind: "tool", + Text: long, + }, + { + Kind: "user", + Text: "next", + }, + { + Kind: "reply", + Text: long, + }, + } + + done := make(chan string) + go func() { + _, stop := run.Run(t.Context(), entries) + done <- stop + }() + + require.Eventually(t, func() bool { + return server.Calls() == 2 + }, 10*time.Second, time.Millisecond) + + run.memory.mutex.Lock() + appended := run.memory.log.Len() + run.memory.mutex.Unlock() + require.Equal(t, 2, appended) + + close(gate) + require.Empty(t, <-done) + require.NoError(t, run.memory.Close()) + + mutex.Lock() + defer mutex.Unlock() + + require.Equal(t, 2, most) + require.Equal(t, 3, server.Calls()) +} + +func TestSlotSummarizerLimitsTheBuildsAtOnce(t *testing.T) { + inner := newFakeSummarizer() + inner.gate = make(chan struct{}) + summarizer := slotSummarizer{ + inner: inner, + slots: make(chan struct{}, 2), + } + + var group sync.WaitGroup + for n := range 6 { + group.Go(func() { + _, err := summarizer.Summarize(t.Context(), Task{ + Level: 0, + Index: n, + Input: []string{"user: " + strings.Repeat("x", 600)}, + }) + require.NoError(t, err) + }) + } + + require.Eventually(t, func() bool { + return inner.Active() == 2 + }, 10*time.Second, time.Millisecond) + + close(inner.gate) + group.Wait() + require.Equal(t, 2, inner.peak) +} + +func TestReadUsageFileSortsByCallStart(t *testing.T) { + path := filepath.Join(t.TempDir(), "usage.jsonl") + usage := usageLog{ + path: path, + } + + for _, date := range []string{"2026-01-01T12:00:02Z", "2026-01-01T12:00:00.5Z", "2026-01-01T12:00:01.25Z"} { + require.NoError(t, usage.Write(usageRecord{ + Date: date, + })) + } + + records, err := readUsageFile(path) + require.NoError(t, err) + require.Equal(t, []string{"2026-01-01T12:00:00.5Z", "2026-01-01T12:00:01.25Z", "2026-01-01T12:00:02Z"}, []string{records[0].Date, records[1].Date, records[2].Date}) +} + +func TestTimeTextCountsTheCallsAtOnce(t *testing.T) { + start := time.Date(2026, 1, 1, 12, 0, 0, 0, time.UTC) + record := func(offset, seconds float64) usageRecord { + return usageRecord{ + Date: start.Add(time.Duration(offset * float64(time.Second))).Format(time.RFC3339Nano), + Seconds: seconds, + } + } + + text := timeText([]usageRecord{record(0, 10), record(2, 4), record(3, 2), record(10, 10)}) + require.Equal(t, "time: 20 s in all, 12.0 calls per minute, 6.5 s per call (median 10.0, max 10.0), calls at once: at most 3, 1.3 on average", text) } // probeLog returns count messages of made-up work on a small bakery order @@ -270,6 +413,22 @@ func (g *probeWriter) pick(words []string) string { return words[g.random.IntN(len(words))] } +func newTestProbeRun(t *testing.T, server *fakeResponses) *probeRun { + t.Helper() + + root := t.TempDir() + usagePath := filepath.Join(root, "usage.jsonl") + client := newClient(server.URL, apiKey(testKey), "onechat-"+probeName, usagePath) + memory, err := openMemory(t.Context(), root, newTestClock().Now, newCompactor(client, probeName, defaultConfig().Compactor)) + require.NoError(t, err) + + return &probeRun{ + memory: memory, + client: client, + usagePath: usagePath, + } +} + var ( people = []string{"Mara", "Joel", "Ines", "Tomas", "Priya", "Lena"} breads = []string{"rye", "sourdough", "baguette", "brioche", "spelt", "seed loaf"}