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

   1 package main
   2
   3 import (
   4         "encoding/json"
   5         "io"
   6         "net/http"
   7         "os"
   8         "path/filepath"
   9         "strings"
  10         "sync"
  11         "testing"
  12         "time"
  13
  14         "github.com/stretchr/testify/require"
  15 )
  16
  17 func TestLoginWritesTheTokenFile(t *testing.T) {
  18         server := newChatGPTFake(t)
  19         server.pending = 2
  20         auth := newTestChatGPTAuth(t, server)
  21
  22         var out strings.Builder
  23
  24         err := auth.Login(t.Context(), &out)
  25         require.NoError(t, err)
  26
  27         require.Contains(t, out.String(), server.URL+"/codex/device")
  28         require.Contains(t, out.String(), "CODE-1234")
  29         require.Equal(t, 3, server.polls)
  30
  31         form := server.exchange
  32         require.Equal(t, "authorization_code", form.Get("grant_type"))
  33         require.Equal(t, "auth-code", form.Get("code"))
  34         require.Equal(t, "verifier", form.Get("code_verifier"))
  35         require.Equal(t, chatgptClientID, form.Get("client_id"))
  36         require.Equal(t, server.URL+"/deviceauth/callback", form.Get("redirect_uri"))
  37
  38         tokens := chatgptReadTokenFile(t, auth.path)
  39         require.Equal(t, "refresh-1", tokens.RefreshToken)
  40         require.Equal(t, chatgptTestAccount, tokens.AccountID)
  41         require.Equal(t, server.access[0], tokens.AccessToken)
  42         require.NotContains(t, out.String(), tokens.AccessToken)
  43         require.NotContains(t, out.String(), tokens.RefreshToken)
  44
  45         info, err := os.Stat(auth.path)
  46         require.NoError(t, err)
  47         require.Equal(t, os.FileMode(0o600), info.Mode().Perm())
  48 }
  49
  50 func TestLoginFailsWhenTheCodeIsDenied(t *testing.T) {
  51         server := newChatGPTFake(t)
  52         server.deny = true
  53         auth := newTestChatGPTAuth(t, server)
  54
  55         err := auth.Login(t.Context(), io.Discard)
  56         require.ErrorContains(t, err, "HTTP 400")
  57         require.NoFileExists(t, auth.path)
  58 }
  59
  60 func TestAuthorizeUsesAFreshTokenWithoutRefresh(t *testing.T) {
  61         server := newChatGPTFake(t)
  62         auth := newTestChatGPTAuth(t, server)
  63         chatgptWriteTestTokens(t, auth, server.token(time.Hour), "refresh-0")
  64
  65         header := http.Header{}
  66         err := auth.Authorize(t.Context(), header)
  67         require.NoError(t, err)
  68         require.Equal(t, "Bearer "+server.access[0], header.Get("Authorization"))
  69         require.Equal(t, chatgptTestAccount, header.Get("ChatGPT-Account-Id"))
  70         require.Equal(t, 0, server.refreshes)
  71 }
  72
  73 func TestAuthorizeRefreshesNearTheExpiry(t *testing.T) {
  74         testCases := []struct {
  75                 name        string
  76                 newRefresh  string
  77                 wantRefresh string
  78         }{
  79                 {
  80                         name:        "new refresh token kept",
  81                         newRefresh:  "refresh-new",
  82                         wantRefresh: "refresh-new",
  83                 },
  84                 {
  85                         name:        "old refresh token kept without a new one",
  86                         newRefresh:  "",
  87                         wantRefresh: "refresh-0",
  88                 },
  89         }
  90
  91         for _, tc := range testCases {
  92                 t.Run(tc.name, func(t *testing.T) {
  93                         server := newChatGPTFake(t)
  94                         server.newRefresh = tc.newRefresh
  95                         auth := newTestChatGPTAuth(t, server)
  96                         chatgptWriteTestTokens(t, auth, server.token(59*time.Second), "refresh-0")
  97
  98                         header := http.Header{}
  99                         err := auth.Authorize(t.Context(), header)
 100                         require.NoError(t, err)
 101                         require.Equal(t, 1, server.refreshes)
 102                         require.Equal(t, "refresh-0", server.refreshed)
 103
 104                         tokens := chatgptReadTokenFile(t, auth.path)
 105                         require.Equal(t, server.access[1], tokens.AccessToken)
 106                         require.Equal(t, tc.wantRefresh, tokens.RefreshToken)
 107                         require.Equal(t, chatgptTestAccount, tokens.AccountID)
 108                         require.Equal(t, "Bearer "+server.access[1], header.Get("Authorization"))
 109
 110                         err = auth.Authorize(t.Context(), header)
 111                         require.NoError(t, err)
 112                         require.Equal(t, 1, server.refreshes)
 113                 })
 114         }
 115 }
 116
 117 func TestAuthorizeRunsOneRefreshAtATime(t *testing.T) {
 118         server := newChatGPTFake(t)
 119         server.newRefresh = "refresh-new"
 120         auth := newTestChatGPTAuth(t, server)
 121         chatgptWriteTestTokens(t, auth, server.token(time.Second), "refresh-0")
 122
 123         var group sync.WaitGroup
 124
 125         errs := make(chan error, 8)
 126         for range 8 {
 127                 group.Go(func() {
 128                         errs <- auth.Authorize(t.Context(), http.Header{})
 129                 })
 130         }
 131
 132         group.Wait()
 133         close(errs)
 134
 135         for err := range errs {
 136                 require.NoError(t, err)
 137         }
 138
 139         require.Equal(t, 1, server.refreshes)
 140 }
 141
 142 func TestAuthorizeNeedsALogin(t *testing.T) {
 143         server := newChatGPTFake(t)
 144         auth := newTestChatGPTAuth(t, server)
 145
 146         err := auth.Authorize(t.Context(), http.Header{})
 147         require.ErrorIs(t, err, errChatGPTLoginNeeded)
 148
 149         server.failRefresh = true
 150         chatgptWriteTestTokens(t, auth, server.token(-time.Minute), "refresh-0")
 151
 152         before, err := os.ReadFile(auth.path)
 153         require.NoError(t, err)
 154
 155         err = auth.Authorize(t.Context(), http.Header{})
 156         require.ErrorIs(t, err, errChatGPTLoginNeeded)
 157         require.NotContains(t, err.Error(), "refresh-0")
 158
 159         after, err := os.ReadFile(auth.path)
 160         require.NoError(t, err)
 161         require.Equal(t, before, after)
 162 }
 163
 164 func TestForcedRefreshLetsTheNextCallWork(t *testing.T) {
 165         server := newChatGPTFake(t)
 166         auth := newTestChatGPTAuth(t, server)
 167         chatgptWriteTestTokens(t, auth, server.token(-time.Hour), "refresh-0")
 168
 169         responses := newModelFake(t)
 170         client := newModelClient(responses.URL, auth, "agent-key", filepath.Join(t.TempDir(), "usage.jsonl"))
 171         _, err := client.Create(t.Context(), modelTestRequest())
 172         require.NoError(t, err)
 173
 174         require.Equal(t, 1, server.refreshes)
 175
 176         call := responses.Call(0)
 177         require.Equal(t, "Bearer "+server.access[1], call.header.Get("Authorization"))
 178         require.Equal(t, chatgptTestAccount, call.header.Get("ChatGPT-Account-Id"))
 179         require.Equal(t, server.access[1], chatgptReadTokenFile(t, auth.path).AccessToken)
 180 }
 181
 182 func TestModelEndpointFollowsTheConfig(t *testing.T) {
 183         dir := t.TempDir()
 184         config := defaultConfig()
 185
 186         t.Setenv("OPENAI_API_KEY", modelTestKey)
 187
 188         address, auth, err := modelEndpoint(config, dir)
 189         require.NoError(t, err)
 190         require.Equal(t, modelChatGPTURL, address)
 191         require.IsType(t, &ChatGPTAuth{}, auth)
 192         require.Equal(t, filepath.Join(dir, chatgptTokenFileName), auth.(*ChatGPTAuth).path)
 193
 194         config.Endpoint = configEndpointAPI
 195         address, auth, err = modelEndpoint(config, dir)
 196         require.NoError(t, err)
 197         require.Equal(t, modelAPIURL, address)
 198         require.Equal(t, ModelAPIKey(modelTestKey), auth)
 199
 200         t.Setenv("OPENAI_API_KEY", "")
 201
 202         _, _, err = modelEndpoint(config, dir)
 203         require.ErrorContains(t, err, "OPENAI_API_KEY")
 204 }
 205
 206 func chatgptReadTokenFile(t *testing.T, path string) ChatGPTTokenFile {
 207         t.Helper()
 208
 209         data, err := os.ReadFile(path)
 210         require.NoError(t, err)
 211
 212         var tokens ChatGPTTokenFile
 213
 214         require.NoError(t, json.Unmarshal(data, &tokens))
 215
 216         return tokens
 217 }