X-Git-Url: https://git.giacomolli.uk/onechat/blobdiff_plain/6d1d9d9a7cf90b25a2e10af1c5bc56faee7e494c..396af2e18eb02d4ec1e35edbc62f01e84bae2b46:/chatgpt_test.go?ds=inline diff --git a/chatgpt_test.go b/chatgpt_test.go index 982a6ff..a644daf 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,9 +14,6 @@ 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 @@ -192,7 +185,7 @@ func TestModelEndpointFollowsTheConfig(t *testing.T) { require.IsType(t, &ChatGPTAuth{}, auth) require.Equal(t, filepath.Join(dir, chatgptTokenFileName), auth.(*ChatGPTAuth).path) - config.Endpoint = endpointAPI + config.Endpoint = configEndpointAPI address, auth, err = modelEndpoint(config, dir) require.NoError(t, err) require.Equal(t, modelAPIURL, address) @@ -203,164 +196,6 @@ func TestModelEndpointFollowsTheConfig(t *testing.T) { 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() @@ -372,23 +207,3 @@ func chatgptReadTokenFile(t *testing.T, path string) ChatGPTTokenFile { 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)) - } -}