X-Git-Url: https://git.giacomolli.uk/onechat/blobdiff_plain/7cc0014244124aea59672904de65a979a8e5fb2c..871743f6ab20d68dcede1673a9764730fda38308:/chatgpt_test.go?ds=inline diff --git a/chatgpt_test.go b/chatgpt_test.go index c6db9bd..cb144fb 100644 --- a/chatgpt_test.go +++ b/chatgpt_test.go @@ -1,13 +1,9 @@ package main import ( - "encoding/base64" "encoding/json" - "fmt" "io" "net/http" - "net/http/httptest" - "net/url" "os" "path/filepath" "strings" @@ -18,15 +14,13 @@ import ( "github.com/stretchr/testify/require" ) -// chatgptTestAccount is the fake ChatGPT account ID in the fake tokens. -const chatgptTestAccount = "account-test" - func TestLoginWritesTheTokenFile(t *testing.T) { server := newChatGPTFake(t) server.pending = 2 auth := newTestChatGPTAuth(t, server) var out strings.Builder + err := auth.Login(t.Context(), &out) require.NoError(t, err) @@ -127,6 +121,7 @@ func TestAuthorizeRunsOneRefreshAtATime(t *testing.T) { chatgptWriteTestTokens(t, auth, server.token(time.Second), "refresh-0") var group sync.WaitGroup + errs := make(chan error, 8) for range 8 { group.Go(func() { @@ -136,6 +131,7 @@ func TestAuthorizeRunsOneRefreshAtATime(t *testing.T) { group.Wait() close(errs) + for err := range errs { require.NoError(t, err) } @@ -152,6 +148,7 @@ func TestAuthorizeNeedsALogin(t *testing.T) { server.failRefresh = true chatgptWriteTestTokens(t, auth, server.token(-time.Minute), "refresh-0") + before, err := os.ReadFile(auth.path) require.NoError(t, err) @@ -175,6 +172,7 @@ func TestForcedRefreshLetsTheNextCallWork(t *testing.T) { require.NoError(t, err) require.Equal(t, 1, server.refreshes) + call := responses.Call(0) require.Equal(t, "Bearer "+server.access[1], call.header.Get("Authorization")) require.Equal(t, chatgptTestAccount, call.header.Get("ChatGPT-Account-Id")) @@ -186,6 +184,7 @@ func TestModelEndpointFollowsTheConfig(t *testing.T) { config := defaultConfig() t.Setenv("OPENAI_API_KEY", modelTestKey) + address, auth, err := modelEndpoint(config, dir) require.NoError(t, err) require.Equal(t, modelChatGPTURL, address) @@ -199,168 +198,11 @@ func TestModelEndpointFollowsTheConfig(t *testing.T) { require.Equal(t, ModelAPIKey(modelTestKey), auth) t.Setenv("OPENAI_API_KEY", "") + _, _, err = modelEndpoint(config, dir) require.ErrorContains(t, err, "OPENAI_API_KEY") } -// ChatGPTFake is an httptest login service. The device token endpoint answers -// 403 pending times before it grants the code. Each token answer has a new -// access token, kept in access. -type ChatGPTFake struct { - *httptest.Server - mutex sync.Mutex - pending int - deny bool - polls int - exchange url.Values - refreshes int - refreshed string - newRefresh string - failRefresh bool - access []string -} - -// token returns a new access token that expires after ttl, for a token file -// that a test writes. -func (f *ChatGPTFake) token(ttl time.Duration) string { - f.mutex.Lock() - defer f.mutex.Unlock() - - return f.newToken(ttl) -} - -// newToken returns a new fake JWT with the expiry and the account claim, and -// keeps it in access. The caller holds the mutex. -func (f *ChatGPTFake) newToken(ttl time.Duration) string { - claims, err := json.Marshal(map[string]any{ - "exp": time.Now().Add(ttl).Unix(), - chatgptAuthClaim: map[string]any{ - "chatgpt_account_id": chatgptTestAccount, - }, - "serial": len(f.access), - }) - if err != nil { - panic(err) - } - - encode := base64.RawURLEncoding.EncodeToString - token := encode([]byte(`{"alg":"none"}`)) + "." + encode(claims) + "." + encode([]byte("fake")) - f.access = append(f.access, token) - - return token -} - -func (f *ChatGPTFake) userCode(w http.ResponseWriter, r *http.Request) { - var body map[string]string - err := json.NewDecoder(r.Body).Decode(&body) - if err != nil || body["client_id"] != chatgptClientID { - http.Error(w, "bad client", http.StatusBadRequest) - return - } - - chatgptWriteJSON(w, map[string]any{ - "device_auth_id": "device-1", - "user_code": "CODE-1234", - "interval": "0", - }) -} - -func (f *ChatGPTFake) deviceToken(w http.ResponseWriter, r *http.Request) { - f.mutex.Lock() - defer f.mutex.Unlock() - - var body map[string]string - err := json.NewDecoder(r.Body).Decode(&body) - if err != nil || body["device_auth_id"] != "device-1" || body["user_code"] != "CODE-1234" { - http.Error(w, "bad device", http.StatusBadRequest) - return - } - - f.polls++ - if f.deny { - http.Error(w, "denied", http.StatusBadRequest) - return - } - - if f.polls <= f.pending { - http.Error(w, "pending", http.StatusForbidden) - return - } - - chatgptWriteJSON(w, map[string]any{ - "authorization_code": "auth-code", - "code_challenge": "challenge", - "code_verifier": "verifier", - }) -} - -func (f *ChatGPTFake) oauthToken(w http.ResponseWriter, r *http.Request) { - f.mutex.Lock() - defer f.mutex.Unlock() - - if r.Header.Get("Content-Type") == "application/x-www-form-urlencoded" { - err := r.ParseForm() - if err != nil { - http.Error(w, err.Error(), http.StatusBadRequest) - return - } - - f.exchange = r.PostForm - chatgptWriteJSON(w, map[string]any{ - "access_token": f.newToken(time.Hour), - "refresh_token": "refresh-1", - "id_token": f.newToken(time.Hour), - }) - - return - } - - var body map[string]string - err := json.NewDecoder(r.Body).Decode(&body) - if err != nil || body["grant_type"] != "refresh_token" || body["client_id"] != chatgptClientID { - http.Error(w, "bad refresh", http.StatusBadRequest) - return - } - - f.refreshes++ - f.refreshed = body["refresh_token"] - if f.failRefresh { - http.Error(w, `{"error": "invalid_grant", "token": "`+body["refresh_token"]+`"}`, http.StatusUnauthorized) - return - } - - answer := map[string]any{ - "access_token": f.newToken(time.Hour), - "id_token": f.newToken(time.Hour), - } - - if f.newRefresh != "" { - answer["refresh_token"] = f.newRefresh - } - - chatgptWriteJSON(w, answer) -} - -func newChatGPTFake(t *testing.T) *ChatGPTFake { - t.Helper() - - f := &ChatGPTFake{} - mux := http.NewServeMux() - mux.HandleFunc("POST /api/accounts/deviceauth/usercode", f.userCode) - mux.HandleFunc("POST /api/accounts/deviceauth/token", f.deviceToken) - mux.HandleFunc("POST /oauth/token", f.oauthToken) - f.Server = httptest.NewServer(mux) - t.Cleanup(f.Close) - - return f -} - -func newTestChatGPTAuth(t *testing.T, server *ChatGPTFake) *ChatGPTAuth { - t.Helper() - - return newChatGPTAuth(filepath.Join(t.TempDir(), chatgptTokenFileName), server.URL) -} - func chatgptReadTokenFile(t *testing.T, path string) ChatGPTTokenFile { t.Helper() @@ -368,27 +210,8 @@ func chatgptReadTokenFile(t *testing.T, path string) ChatGPTTokenFile { require.NoError(t, err) var tokens ChatGPTTokenFile + require.NoError(t, json.Unmarshal(data, &tokens)) return tokens } - -func chatgptWriteTestTokens(t *testing.T, auth *ChatGPTAuth, access, refresh string) { - t.Helper() - - err := auth.write(ChatGPTTokenFile{ - AccessToken: access, - RefreshToken: refresh, - IDToken: access, - AccountID: chatgptTestAccount, - }) - require.NoError(t, err) -} - -func chatgptWriteJSON(w http.ResponseWriter, value any) { - w.Header().Set("Content-Type", "application/json") - err := json.NewEncoder(w).Encode(value) - if err != nil { - panic(fmt.Sprint("encode answer: ", err)) - } -}