Wait for the WhatsApp catch-up before a turn

[onechat] / model_test.go

diff --git a/model_test.go b/model_test.go
index b90eaddca24bec8012cee558e4194585cc6942e2..11a7d00bcddf7ad8f68373675a8428111fa81fb4 100644 (file)
--- a/model_test.go
+++ b/model_test.go
@@ -1,24 +1,16 @@
 package main
 
 import (
 package main
 
 import (

-       "bufio"
-       "encoding/json"
-       "fmt"

        "io"
        "net/http"
        "io"
        "net/http"

-       "net/http/httptest"

        "os"
        "path/filepath"
        "os"
        "path/filepath"

-       "sync"

        "testing"
        "time"
 
        "github.com/stretchr/testify/require"
 )
 
        "testing"
        "time"
 
        "github.com/stretchr/testify/require"
 )
 

-// modelTestKey is a fake API key.
-const modelTestKey = "sk-test-not-a-real-key"
-

 func TestClientSendsTheFixedRequestFields(t *testing.T) {
        server := newModelFake(t)
        client := newTestModelClient(t, server)
 func TestClientSendsTheFixedRequestFields(t *testing.T) {
        server := newModelFake(t)
        client := newTestModelClient(t, server)


@@ -62,9 +54,9 @@ func TestClientSendsTheFixedRequestFields(t *testing.T) {
 }
 
 func TestClientSendsTheSessionHeaderToChatGPT(t *testing.T) {
 }
 
 func TestClientSendsTheSessionHeaderToChatGPT(t *testing.T) {

-       login := newFakeAuth(t)
-       auth := newTestAuth(t, login)
-       writeTestTokens(t, auth, login.token(time.Hour), "refresh-0")
+       login := newChatGPTFake(t)
+       auth := newTestChatGPTAuth(t, login)
+       chatgptWriteTestTokens(t, auth, login.token(time.Hour), "refresh-0")

 
        testCases := []struct {
                name string
 
        testCases := []struct {
                name string


@@ -93,6 +85,19 @@ func TestClientSendsTheSessionHeaderToChatGPT(t *testing.T) {
                        call := server.Call(0)
                        require.Equal(t, tc.want, call.header.Get("session_id"))
                        require.Equal(t, "agent-key", call.body["prompt_cache_key"])
                        call := server.Call(0)
                        require.Equal(t, tc.want, call.header.Get("session_id"))
                        require.Equal(t, "agent-key", call.body["prompt_cache_key"])

+
+                       // A subagent has its own key.
+                       request := modelTestRequest()
+                       request.Agent = "GreenOtter"
+                       _, err = client.Create(t.Context(), request)
+                       require.NoError(t, err)
+
+                       call = server.Call(1)
+                       if tc.want != "" {
+                               require.Equal(t, "agent-key-GreenOtter", call.header.Get("session_id"))
+                       }
+
+                       require.Equal(t, "agent-key-GreenOtter", call.body["prompt_cache_key"])

                })
        }
 }
                })
        }
 }


@@ -152,11 +157,13 @@ func TestClientRecordsEachCall(t *testing.T) {
 
        client := newTestModelClient(t, server)
        request := modelTestRequest()
 
        client := newTestModelClient(t, server)
        request := modelTestRequest()

-       request.Node = &nodeKey{2, 5}
+       request.Node = &NodeKey{2, 5}
+       request.Agent = "GreenOtter"

        _, err := client.Create(t.Context(), request)
        require.NoError(t, err)
        require.NotContains(t, server.Call(0).body, "node")
        require.NotContains(t, string(server.Call(0).raw), "[2,5]")
        _, err := client.Create(t.Context(), request)
        require.NoError(t, err)
        require.NotContains(t, server.Call(0).body, "node")
        require.NotContains(t, string(server.Call(0).raw), "[2,5]")

+       require.NotContains(t, server.Call(0).body, "agent")

 
        _, err = client.Create(t.Context(), modelTestRequest())
        require.ErrorContains(t, err, "HTTP 429")
 
        _, err = client.Create(t.Context(), modelTestRequest())
        require.ErrorContains(t, err, "HTTP 429")


@@ -176,12 +183,14 @@ func TestClientRecordsEachCall(t *testing.T) {
        require.Equal(t, 7, records[0].OutputTokens)
        require.Equal(t, quota, records[0].Quota)
        require.Equal(t, &[2]int{2, 5}, records[0].Node)
        require.Equal(t, 7, records[0].OutputTokens)
        require.Equal(t, quota, records[0].Quota)
        require.Equal(t, &[2]int{2, 5}, records[0].Node)

+       require.Equal(t, "GreenOtter", records[0].Agent)

        require.Equal(t, len("ok"), records[0].TextBytes)
        require.Greater(t, records[0].Seconds, 0.0)
        require.Empty(t, records[0].Error)
        require.NotEmpty(t, records[0].Date)
 
        require.Nil(t, records[1].Node)
        require.Equal(t, len("ok"), records[0].TextBytes)
        require.Greater(t, records[0].Seconds, 0.0)
        require.Empty(t, records[0].Error)
        require.NotEmpty(t, records[0].Date)
 
        require.Nil(t, records[1].Node)

+       require.Empty(t, records[1].Agent)

        require.Zero(t, records[1].TextBytes)
 
        require.Equal(t, http.StatusTooManyRequests, records[1].Status)
        require.Zero(t, records[1].TextBytes)
 
        require.Equal(t, http.StatusTooManyRequests, records[1].Status)


@@ -261,177 +270,8 @@ func TestClientFailsWhenTheResponseFails(t *testing.T) {
        }
 }
 
        }
 }
 

-// ModelFakeCall is one call to a fake endpoint, numbered from 0.
-type ModelFakeCall struct {
-       number int
-       header http.Header
-       raw    []byte
-       body   map[string]any
-}
-
-// ModelFake is an httptest Responses endpoint. It keeps each call, and reply
-// writes the response. The default reply is one short message.
-type ModelFake struct {
-       *httptest.Server
-       mutex sync.Mutex
-       calls []ModelFakeCall
-       reply func(http.ResponseWriter, ModelFakeCall)
-}
-
-// Call returns call number n.
-func (f *ModelFake) Call(n int) ModelFakeCall {
-       f.mutex.Lock()
-       defer f.mutex.Unlock()
-
-       return f.calls[n]
-}
-
-// Calls returns the number of calls.
-func (f *ModelFake) Calls() int {
-       f.mutex.Lock()
-       defer f.mutex.Unlock()
-
-       return len(f.calls)
-}
-
-func (f *ModelFake) serve(w http.ResponseWriter, r *http.Request) {
-       raw, err := io.ReadAll(r.Body)
-       if err != nil {
-               http.Error(w, err.Error(), http.StatusBadRequest)
-               return
-       }
-
-       call := ModelFakeCall{
-               header: r.Header.Clone(),
-               raw:    raw,
-       }
-
-       err = json.Unmarshal(raw, &call.body)
-       if err != nil {
-               http.Error(w, err.Error(), http.StatusBadRequest)
-               return
-       }
-
-       f.mutex.Lock()
-       call.number = len(f.calls)
-       f.calls = append(f.calls, call)
-       reply := f.reply
-       f.mutex.Unlock()
-
-       reply(w, call)
-}
-
-func newModelFake(t *testing.T) *ModelFake {
-       t.Helper()
-
-       f := &ModelFake{
-               reply: func(w http.ResponseWriter, call ModelFakeCall) {
-                       modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("a line")}, 100, 0, 5))
-               },
-       }
-
-       f.Server = httptest.NewServer(http.HandlerFunc(f.serve))
-       t.Cleanup(f.Close)
-
-       return f
-}
-

 func newTestModelClient(t *testing.T, server *ModelFake) *ModelClient {
        t.Helper()
 
        return newModelClient(server.URL, ModelAPIKey(modelTestKey), "agent-key", filepath.Join(t.TempDir(), "usage.jsonl"))
 }
 func newTestModelClient(t *testing.T, server *ModelFake) *ModelClient {
        t.Helper()
 
        return newModelClient(server.URL, ModelAPIKey(modelTestKey), "agent-key", filepath.Join(t.TempDir(), "usage.jsonl"))
 }

-
-func modelTestRequest() ModelRequest {
-       return ModelRequest{
-               Model:  "gpt-test",
-               Effort: "low",
-               Input:  []any{modelUserMessage("hello")},
-       }
-}
-
-// modelWriteEvents writes events as a server-sent event stream.
-func modelWriteEvents(w http.ResponseWriter, events ...any) {
-       w.Header().Set("Content-Type", "text/event-stream")
-       for _, event := range events {
-               data, err := json.Marshal(event)
-               if err != nil {
-                       panic(err)
-               }
-
-               _, err = fmt.Fprintf(w, "event: %s\ndata: %s\n\n", event.(map[string]any)["type"], data)
-               if err != nil {
-                       panic(err)
-               }
-       }
-}
-
-func modelMessageItem(text string) map[string]any {
-       return map[string]any{
-               "type": "message",
-               "role": "assistant",
-               "content": []any{
-                       map[string]any{"type": "output_text", "text": text},
-               },
-       }
-}
-
-func modelCompletedEvent(output []any, input, cached, out int) map[string]any {
-       if output == nil {
-               output = []any{}
-       }
-
-       return map[string]any{
-               "type": "response.completed",
-               "response": map[string]any{
-                       "output": output,
-                       "usage": map[string]any{
-                               "input_tokens":         input,
-                               "input_tokens_details": map[string]any{"cached_tokens": cached},
-                               "output_tokens":        out,
-                       },
-               },
-       }
-}
-
-func modelReadUsage(t *testing.T, path string) []ModelUsageRecord {
-       t.Helper()
-
-       file, err := os.Open(path)
-       require.NoError(t, err)
-
-       defer func() {
-               require.NoError(t, file.Close())
-       }()
-
-       var records []ModelUsageRecord
-       scanner := bufio.NewScanner(file)
-       for scanner.Scan() {
-               var record ModelUsageRecord
-               require.NoError(t, json.Unmarshal(scanner.Bytes(), &record))
-               records = append(records, record)
-       }
-
-       require.NoError(t, scanner.Err())
-
-       return records
-}
-
-func modelReasoningItem(id, encrypted string) map[string]any {
-       return map[string]any{
-               "type":              "reasoning",
-               "id":                id,
-               "encrypted_content": encrypted,
-               "summary":           []any{},
-       }
-}
-
-func modelFunctionCallItem(id, callID, name, arguments string) map[string]any {
-       return map[string]any{
-               "type":      "function_call",
-               "id":        id,
-               "call_id":   callID,
-               "name":      name,
-               "arguments": arguments,
-       }
-}