X-Git-Url: https://git.giacomolli.uk/onechat/blobdiff_plain/6b21fcf78372995c6d8130c63f2f8aa2609bf7fe..98d137c3e945a09f76e85bd18238558c668fdca4:/chatgpt_fake_helper_test.go?ds=sidebyside

diff --git a/chatgpt_fake_helper_test.go b/chatgpt_fake_helper_test.go
index dfcf17d..a25ed12 100644
--- a/chatgpt_fake_helper_test.go
+++ b/chatgpt_fake_helper_test.go
@@ -35,41 +35,34 @@ type ChatGPTFake struct {
 	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()
+func newChatGPTFake(t *testing.T) *ChatGPTFake {
+	t.Helper()
 
-	return f.newToken(ttl)
-}
+	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)
 
-// 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)
-	}
+	f.Server = httptest.NewServer(mux)
+	t.Cleanup(f.Close)
 
-	encode := base64.RawURLEncoding.EncodeToString
-	token := encode([]byte(`{"alg":"none"}`)) + "." + encode(claims) + "." + encode([]byte("fake"))
-	f.access = append(f.access, token)
+	return f
+}
 
-	return token
+func newTestChatGPTAuth(t *testing.T, server *ChatGPTFake) *ChatGPTAuth {
+	t.Helper()
+
+	return newChatGPTAuth(filepath.Join(t.TempDir(), chatgptTokenFileName), server.URL)
 }
 
 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
 	}
 
@@ -85,20 +78,24 @@ func (f *ChatGPTFake) deviceToken(w http.ResponseWriter, r *http.Request) {
 	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
 	}
 
@@ -117,6 +114,7 @@ func (f *ChatGPTFake) oauthToken(w http.ResponseWriter, r *http.Request) {
 		err := r.ParseForm()
 		if err != nil {
 			http.Error(w, err.Error(), http.StatusBadRequest)
+
 			return
 		}
 
@@ -131,9 +129,11 @@ func (f *ChatGPTFake) oauthToken(w http.ResponseWriter, r *http.Request) {
 	}
 
 	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
 	}
 
@@ -141,6 +141,7 @@ func (f *ChatGPTFake) oauthToken(w http.ResponseWriter, r *http.Request) {
 	f.refreshed = body["refresh_token"]
 	if f.failRefresh {
 		http.Error(w, `{"error": "invalid_grant", "token": "`+body["refresh_token"]+`"}`, http.StatusUnauthorized)
+
 		return
 	}
 
@@ -156,24 +157,43 @@ func (f *ChatGPTFake) oauthToken(w http.ResponseWriter, r *http.Request) {
 	chatgptWriteJSON(w, answer)
 }
 
-func newChatGPTFake(t *testing.T) *ChatGPTFake {
-	t.Helper()
+// 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)
+	}
 
-	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)
+	encode := base64.RawURLEncoding.EncodeToString
+	token := encode([]byte(`{"alg":"none"}`)) + "." + encode(claims) + "." + encode([]byte("fake"))
+	f.access = append(f.access, token)
 
-	return f
+	return token
 }
 
-func newTestChatGPTAuth(t *testing.T, server *ChatGPTFake) *ChatGPTAuth {
-	t.Helper()
+// 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 newChatGPTAuth(filepath.Join(t.TempDir(), chatgptTokenFileName), server.URL)
+	return f.newToken(ttl)
+}
+
+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))
+	}
 }
 
 func chatgptWriteTestTokens(t *testing.T, auth *ChatGPTAuth, access, refresh string) {
@@ -187,11 +207,3 @@ func chatgptWriteTestTokens(t *testing.T, auth *ChatGPTAuth, access, refresh str
 	})
 	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))
-	}
-}