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 func newModelFake(t *testing.T) *ModelFake {
  38         t.Helper()
  39
  40         f := &ModelFake{
  41                 reply: func(w http.ResponseWriter, _ ModelFakeCall) {
  42                         modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("a line")}, 100, 0, 5))
  43                 },
  44         }
  45
  46         f.Server = httptest.NewServer(http.HandlerFunc(f.serve))
  47         t.Cleanup(f.Close)
  48
  49         return f
  50 }
  51
  52 func (f *ModelFake) serve(w http.ResponseWriter, r *http.Request) {
  53         raw, err := io.ReadAll(r.Body)
  54         if err != nil {
  55                 http.Error(w, err.Error(), http.StatusBadRequest)
  56
  57                 return
  58         }
  59
  60         call := ModelFakeCall{
  61                 header: r.Header.Clone(),
  62                 raw:    raw,
  63         }
  64
  65         err = json.Unmarshal(raw, &call.body)
  66         if err != nil {
  67                 http.Error(w, err.Error(), http.StatusBadRequest)
  68
  69                 return
  70         }
  71
  72         f.mutex.Lock()
  73
  74         call.number = len(f.calls)
  75         f.calls = append(f.calls, call)
  76         reply := f.reply
  77         f.mutex.Unlock()
  78
  79         reply(w, call)
  80 }
  81
  82 // Call returns call number n.
  83 func (f *ModelFake) Call(n int) ModelFakeCall {
  84         f.mutex.Lock()
  85         defer f.mutex.Unlock()
  86
  87         return f.calls[n]
  88 }
  89
  90 // Calls returns the number of calls.
  91 func (f *ModelFake) Calls() int {
  92         f.mutex.Lock()
  93         defer f.mutex.Unlock()
  94
  95         return len(f.calls)
  96 }
  97
  98 // modelWriteEvents writes events as a server-sent event stream.
  99 func modelWriteEvents(w http.ResponseWriter, events ...any) {
 100         w.Header().Set("Content-Type", "text/event-stream")
 101
 102         for _, event := range events {
 103                 data, err := json.Marshal(event)
 104                 if err != nil {
 105                         panic(err)
 106                 }
 107
 108                 _, err = fmt.Fprintf(w, "event: %s\ndata: %s\n\n", event.(map[string]any)["type"], data)
 109                 if err != nil {
 110                         panic(err)
 111                 }
 112         }
 113 }
 114
 115 func modelCompletedEvent(output []any, input, cached, out int) map[string]any {
 116         if output == nil {
 117                 output = []any{}
 118         }
 119
 120         return map[string]any{
 121                 "type": "response.completed",
 122                 "response": map[string]any{
 123                         "output": output,
 124                         "usage": map[string]any{
 125                                 "input_tokens":         input,
 126                                 "input_tokens_details": map[string]any{"cached_tokens": cached},
 127                                 "output_tokens":        out,
 128                         },
 129                 },
 130         }
 131 }
 132
 133 func modelMessageItem(text string) map[string]any {
 134         return map[string]any{
 135                 "type": "message",
 136                 "role": "assistant",
 137                 "content": []any{
 138                         map[string]any{"type": "output_text", "text": text},
 139                 },
 140         }
 141 }
 142
 143 func modelTestRequest() ModelRequest {
 144         return ModelRequest{
 145                 Model:  "gpt-test",
 146                 Effort: "low",
 147                 Input:  []any{modelUserMessage("hello")},
 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 modelFunctionCallItem(id, callID, name, arguments string) map[string]any {
 178         return map[string]any{
 179                 "type":      "function_call",
 180                 "id":        id,
 181                 "call_id":   callID,
 182                 "name":      name,
 183                 "arguments": arguments,
 184         }
 185 }