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