git.giacomolli.uk - onechat/blob - model_fake_helper_test.go

   1 package main
   2
   3 import (
   4         "bufio"
   5         "encoding/json"
   6         "fmt"
   7         "io"
   8         "net/http"
   9         "net/http/httptest"
  10         "os"
  11         "sync"
  12         "testing"
  13
  14         "github.com/stretchr/testify/require"
  15 )
  16
  17 // modelTestKey is a fake API key.
  18 const modelTestKey = "sk-test-not-a-real-key"
  19
  20 // ModelFakeCall is one call to a fake endpoint, numbered from 0.
  21 type ModelFakeCall struct {
  22         number int
  23         header http.Header
  24         raw    []byte
  25         body   map[string]any
  26 }
  27
  28 // ModelFake is an httptest Responses endpoint. It keeps each call, and reply
  29 // writes the response. The default reply is one short message.
  30 type ModelFake struct {
  31         *httptest.Server
  32         mutex sync.Mutex
  33         calls []ModelFakeCall
  34         reply func(http.ResponseWriter, ModelFakeCall)
  35 }
  36
  37 // Call returns call number n.
  38 func (f *ModelFake) Call(n int) ModelFakeCall {
  39         f.mutex.Lock()
  40         defer f.mutex.Unlock()
  41
  42         return f.calls[n]
  43 }
  44
  45 // Calls returns the number of calls.
  46 func (f *ModelFake) Calls() int {
  47         f.mutex.Lock()
  48         defer f.mutex.Unlock()
  49
  50         return len(f.calls)
  51 }
  52
  53 func (f *ModelFake) serve(w http.ResponseWriter, r *http.Request) {
  54         raw, err := io.ReadAll(r.Body)
  55         if err != nil {
  56                 http.Error(w, err.Error(), http.StatusBadRequest)
  57
  58                 return
  59         }
  60
  61         call := ModelFakeCall{
  62                 header: r.Header.Clone(),
  63                 raw:    raw,
  64         }
  65
  66         err = json.Unmarshal(raw, &call.body)
  67         if err != nil {
  68                 http.Error(w, err.Error(), http.StatusBadRequest)
  69
  70                 return
  71         }
  72
  73         f.mutex.Lock()
  74
  75         call.number = len(f.calls)
  76         f.calls = append(f.calls, call)
  77         reply := f.reply
  78         f.mutex.Unlock()
  79
  80         reply(w, call)
  81 }
  82
  83 func newModelFake(t *testing.T) *ModelFake {
  84         t.Helper()
  85
  86         f := &ModelFake{
  87                 reply: func(w http.ResponseWriter, _ ModelFakeCall) {
  88                         modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("a line")}, 100, 0, 5))
  89                 },
  90         }
  91
  92         f.Server = httptest.NewServer(http.HandlerFunc(f.serve))
  93         t.Cleanup(f.Close)
  94
  95         return f
  96 }
  97
  98 func modelTestRequest() ModelRequest {
  99         return ModelRequest{
 100                 Model:  "gpt-test",
 101                 Effort: "low",
 102                 Input:  []any{modelUserMessage("hello")},
 103         }
 104 }
 105
 106 // modelWriteEvents writes events as a server-sent event stream.
 107 func modelWriteEvents(w http.ResponseWriter, events ...any) {
 108         w.Header().Set("Content-Type", "text/event-stream")
 109
 110         for _, event := range events {
 111                 data, err := json.Marshal(event)
 112                 if err != nil {
 113                         panic(err)
 114                 }
 115
 116                 _, err = fmt.Fprintf(w, "event: %s\ndata: %s\n\n", event.(map[string]any)["type"], data)
 117                 if err != nil {
 118                         panic(err)
 119                 }
 120         }
 121 }
 122
 123 func modelMessageItem(text string) map[string]any {
 124         return map[string]any{
 125                 "type": "message",
 126                 "role": "assistant",
 127                 "content": []any{
 128                         map[string]any{"type": "output_text", "text": text},
 129                 },
 130         }
 131 }
 132
 133 func modelCompletedEvent(output []any, input, cached, out int) map[string]any {
 134         if output == nil {
 135                 output = []any{}
 136         }
 137
 138         return map[string]any{
 139                 "type": "response.completed",
 140                 "response": map[string]any{
 141                         "output": output,
 142                         "usage": map[string]any{
 143                                 "input_tokens":         input,
 144                                 "input_tokens_details": map[string]any{"cached_tokens": cached},
 145                                 "output_tokens":        out,
 146                         },
 147                 },
 148         }
 149 }
 150
 151 func modelReadUsage(t *testing.T, path string) []ModelUsageRecord {
 152         t.Helper()
 153
 154         file, err := os.Open(path)
 155         require.NoError(t, err)
 156
 157         defer func() {
 158                 require.NoError(t, file.Close())
 159         }()
 160
 161         var records []ModelUsageRecord
 162
 163         scanner := bufio.NewScanner(file)
 164         for scanner.Scan() {
 165                 var record ModelUsageRecord
 166
 167                 require.NoError(t, json.Unmarshal(scanner.Bytes(), &record))
 168
 169                 records = append(records, record)
 170         }
 171
 172         require.NoError(t, scanner.Err())
 173
 174         return records
 175 }
 176
 177 func modelReasoningItem(id, encrypted string) map[string]any {
 178         return map[string]any{
 179                 "type":              "reasoning",
 180                 "id":                id,
 181                 "encrypted_content": encrypted,
 182                 "summary":           []any{},
 183         }
 184 }
 185
 186 func modelFunctionCallItem(id, callID, name, arguments string) map[string]any {
 187         return map[string]any{
 188                 "type":      "function_call",
 189                 "id":        id,
 190                 "call_id":   callID,
 191                 "name":      name,
 192                 "arguments": arguments,
 193         }
 194 }