X-Git-Url: https://git.giacomolli.uk/onechat/blobdiff_plain/fe39a670d7af062af089bdba4dd07c3f0aee4048..153f3f9b4b09d521412eacca30eddeed395c3f41:/probe_test.go

diff --git a/probe_test.go b/probe_test.go
index 0651281..afc6e2f 100644
--- a/probe_test.go
+++ b/probe_test.go
@@ -70,8 +70,8 @@ func TestProbeLogHasTheSyntheticMix(t *testing.T) {
 }
 
 func TestProbeStopRules(t *testing.T) {
-	quota := func(primary, secondary string) usageRecord {
-		return usageRecord{
+	quota := func(primary, secondary string) ModelUsageRecord {
+		return ModelUsageRecord{
 			Quota: map[string]string{
 				"x-codex-primary-used-percent":   primary,
 				"x-codex-secondary-used-percent": secondary,
@@ -82,7 +82,7 @@ func TestProbeStopRules(t *testing.T) {
 
 	testCases := []struct {
 		name    string
-		records []usageRecord
+		records []ModelUsageRecord
 		want    string
 	}{
 		{
@@ -92,17 +92,17 @@ func TestProbeStopRules(t *testing.T) {
 		},
 		{
 			name:    "growth within the limit",
-			records: []usageRecord{quota("10", "40"), quota("25", "41")},
+			records: []ModelUsageRecord{quota("10", "40"), quota("25", "41")},
 			want:    "",
 		},
 		{
 			name:    "growth over the limit",
-			records: []usageRecord{quota("10", "40"), quota("12", "55.5")},
+			records: []ModelUsageRecord{quota("10", "40"), quota("12", "55.5")},
 			want:    "x-codex-secondary-used-percent grew from 40 to 55.5",
 		},
 		{
 			name: "failed call",
-			records: []usageRecord{quota("10", "40"), {
+			records: []ModelUsageRecord{quota("10", "40"), {
 				Status: 400,
 				Error:  "model call: HTTP 400",
 			}},
@@ -118,10 +118,10 @@ func TestProbeStopRules(t *testing.T) {
 }
 
 func TestProbeRunStopsWhenTheQuotaGrows(t *testing.T) {
-	server := newFakeResponses(t)
-	server.reply = func(w http.ResponseWriter, call fakeCall) {
+	server := newModelFake(t)
+	server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
 		w.Header().Set("X-Codex-Primary-Used-Percent", strconv.Itoa(10+call.number))
-		writeEvents(w, completedEvent([]any{messageItem(strings.Repeat("s", 400))}, 1000, 800, 50))
+		modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem(strings.Repeat("s", 400))}, 1000, 800, 50))
 	}
 
 	run := newTestProbeRun(t, server)
@@ -148,7 +148,7 @@ func TestProbeRunStopsWhenTheQuotaGrows(t *testing.T) {
 }
 
 func TestProbeRunStopsAtTheCallLimit(t *testing.T) {
-	server := newFakeResponses(t)
+	server := newModelFake(t)
 	run := newTestProbeRun(t, server)
 	run.client.LimitCalls(5)
 
@@ -170,8 +170,8 @@ 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) {
+	server := newModelFake(t)
+	server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
 		mutex.Lock()
 		running++
 		most = max(most, running)
@@ -183,7 +183,7 @@ func TestProbeRunWaitsOnlyBeforeUserMessages(t *testing.T) {
 		running--
 		mutex.Unlock()
 
-		writeEvents(w, completedEvent([]any{messageItem("a line")}, 100, 0, 5))
+		modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("a line")}, 100, 0, 5))
 	}
 
 	run := newTestProbeRun(t, server)
@@ -264,12 +264,12 @@ func TestSlotSummarizerLimitsTheBuildsAtOnce(t *testing.T) {
 
 func TestReadUsageFileSortsByCallStart(t *testing.T) {
 	path := filepath.Join(t.TempDir(), "usage.jsonl")
-	usage := usageLog{
+	usage := ModelUsageLog{
 		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{
+		require.NoError(t, usage.Write(ModelUsageRecord{
 			Date: date,
 		}))
 	}
@@ -281,14 +281,14 @@ func TestReadUsageFileSortsByCallStart(t *testing.T) {
 
 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{
+	record := func(offset, seconds float64) ModelUsageRecord {
+		return ModelUsageRecord{
 			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)})
+	text := timeText([]ModelUsageRecord{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)
 }
 
@@ -413,13 +413,13 @@ func (g *probeWriter) pick(words []string) string {
 	return words[g.random.IntN(len(words))]
 }
 
-func newTestProbeRun(t *testing.T, server *fakeResponses) *probeRun {
+func newTestProbeRun(t *testing.T, server *ModelFake) *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))
+	client := newModelClient(server.URL, ModelAPIKey(modelTestKey), "onechat-"+probeName, usagePath)
+	memory, err := openMemory(t.Context(), root, newTestClock().Now, newCompactor(client, probePrompt, defaultConfig().Compactor))
 	require.NoError(t, err)
 
 	return &probeRun{