78 lines
3.3 KiB
Go
78 lines
3.3 KiB
Go
|
package observation
|
||
|
|
|
||
|
|
import (
|
||
|
|
"fmt"
|
||
|
|
"sync"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"go.maximumdirect.net/promptkit/internal/domain"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestGroupTerminalAccounting(t *testing.T) {
|
||
|
|
var group Group
|
||
|
|
empty := group.Snapshot()
|
||
|
|
if !empty.TerminalCoverageComplete || !empty.Usage.PromptTokens.Known || !empty.Usage.PromptTokens.Complete || empty.Usage.PromptTokens.Value != 0 {
|
||
|
|
t.Fatalf("empty = %#v", empty)
|
||
|
|
}
|
||
|
|
group.Observe(domain.ExecutionEvent{Kind: domain.ExecutionAttempt, RunID: "pending", Attempts: 99})
|
||
|
|
pending := group.Snapshot()
|
||
|
|
if pending.TerminalCoverageComplete || pending.Attempts != 0 || pending.Usage.PromptTokens.Known || pending.Usage.PromptTokens.Complete {
|
||
|
|
t.Fatalf("pending = %#v", pending)
|
||
|
|
}
|
||
|
|
zero := domain.ObservedTokenCount{Known: true, Complete: true}
|
||
|
|
known := domain.ObservedTokenCount{Value: 4, Known: true, Complete: true}
|
||
|
|
partial := domain.ObservedTokenCount{Value: 2, Known: true, Complete: false}
|
||
|
|
group.Observe(domain.ExecutionEvent{Kind: domain.ExecutionFinished, RunID: "first", Outcome: domain.ExecutionOutcomeError, Attempts: 2, TransientRetries: 1, RepairCalls: 1, Usage: domain.ObservedTokenUsage{PromptTokens: known, CompletionTokens: zero, TotalTokens: partial, CachedTokens: zero}})
|
||
|
|
group.Observe(domain.ExecutionEvent{Kind: domain.ExecutionFinished, RunID: "first", Attempts: 100})
|
||
|
|
group.Observe(domain.ExecutionEvent{Kind: domain.ExecutionStarted, RunID: "first"})
|
||
|
|
group.Observe(domain.ExecutionEvent{Kind: domain.ExecutionFinished, RunID: "pending", Attempts: 3, Usage: domain.ObservedTokenUsage{PromptTokens: partial, CompletionTokens: zero, TotalTokens: domain.ObservedTokenCount{}, CachedTokens: zero}})
|
||
|
|
got := group.Snapshot()
|
||
|
|
if !got.TerminalCoverageComplete || len(got.Members) != 2 || got.Attempts != 5 || got.TransientRetries != 1 || got.RepairCalls != 1 {
|
||
|
|
t.Fatalf("summary = %#v", got)
|
||
|
|
}
|
||
|
|
if got.Usage.PromptTokens != (domain.ObservedTokenCount{Value: 6, Known: true}) || got.Usage.CompletionTokens != zero || got.Usage.TotalTokens != partial || got.Usage.CachedTokens != zero || got.Usage.CacheWriteTokens != (domain.ObservedTokenCount{}) {
|
||
|
|
t.Fatalf("usage = %#v", got.Usage)
|
||
|
|
}
|
||
|
|
got.Members[0].RunID = "changed"
|
||
|
|
for _, member := range group.Snapshot().Members {
|
||
|
|
if member.RunID == "changed" || !member.TerminalObserved {
|
||
|
|
t.Fatalf("member = %#v", member)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestGroupConcurrentSnapshots(t *testing.T) {
|
||
|
|
var group Group
|
||
|
|
var wg sync.WaitGroup
|
||
|
|
const count = 100
|
||
|
|
for i := 0; i < count; i++ {
|
||
|
|
wg.Add(1)
|
||
|
|
go func(i int) {
|
||
|
|
defer wg.Done()
|
||
|
|
id := fmt.Sprint(i)
|
||
|
|
group.Observe(domain.ExecutionEvent{RunID: id, Kind: domain.ExecutionStarted})
|
||
|
|
group.Observe(domain.ExecutionEvent{RunID: id, Kind: domain.ExecutionFinished, Attempts: 1})
|
||
|
|
}(i)
|
||
|
|
}
|
||
|
|
for i := 0; i < count; i++ {
|
||
|
|
snapshot := group.Snapshot()
|
||
|
|
seen, terminals := map[string]bool{}, 0
|
||
|
|
for _, member := range snapshot.Members {
|
||
|
|
if seen[member.RunID] {
|
||
|
|
t.Fatal("duplicate member")
|
||
|
|
}
|
||
|
|
seen[member.RunID] = true
|
||
|
|
if member.TerminalObserved {
|
||
|
|
terminals++
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if snapshot.Attempts != uint64(terminals) || snapshot.TerminalCoverageComplete != (terminals == len(snapshot.Members)) {
|
||
|
|
t.Fatalf("inconsistent snapshot = %#v", snapshot)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
wg.Wait()
|
||
|
|
if final := group.Snapshot(); len(final.Members) != count || final.Attempts != count || !final.TerminalCoverageComplete {
|
||
|
|
t.Fatalf("final = %#v", final)
|
||
|
|
}
|
||
|
|
}
|