X-Git-Url: http://git.giacomolli.uk/onechat/blobdiff_plain/02cdaff8be27f37f44db8892bb12554f2d702323:/auth_test.go..6d1d9d9a7cf90b25a2e10af1c5bc56faee7e494c:/chatgpt_test.go?ds=inline

diff --git a/auth_test.go b/chatgpt_test.go
similarity index 72%
rename from auth_test.go
rename to chatgpt_test.go
index 0621605..982a6ff 100644
--- a/auth_test.go
+++ b/chatgpt_test.go
@@ -18,13 +18,13 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-// testAccount is the fake ChatGPT account ID in the fake tokens.
-const testAccount = "account-test"
+// chatgptTestAccount is the fake ChatGPT account ID in the fake tokens.
+const chatgptTestAccount = "account-test"
 
 func TestLoginWritesTheTokenFile(t *testing.T) {
-	server := newFakeAuth(t)
+	server := newChatGPTFake(t)
 	server.pending = 2
-	auth := newTestAuth(t, server)
+	auth := newTestChatGPTAuth(t, server)
 
 	var out strings.Builder
 	err := auth.Login(t.Context(), &out)
@@ -38,12 +38,12 @@ func TestLoginWritesTheTokenFile(t *testing.T) {
 	require.Equal(t, "authorization_code", form.Get("grant_type"))
 	require.Equal(t, "auth-code", form.Get("code"))
 	require.Equal(t, "verifier", form.Get("code_verifier"))
-	require.Equal(t, codexClientID, form.Get("client_id"))
+	require.Equal(t, chatgptClientID, form.Get("client_id"))
 	require.Equal(t, server.URL+"/deviceauth/callback", form.Get("redirect_uri"))
 
-	tokens := readTokenFile(t, auth.path)
+	tokens := chatgptReadTokenFile(t, auth.path)
 	require.Equal(t, "refresh-1", tokens.RefreshToken)
-	require.Equal(t, testAccount, tokens.AccountID)
+	require.Equal(t, chatgptTestAccount, tokens.AccountID)
 	require.Equal(t, server.access[0], tokens.AccessToken)
 	require.NotContains(t, out.String(), tokens.AccessToken)
 	require.NotContains(t, out.String(), tokens.RefreshToken)
@@ -54,9 +54,9 @@ func TestLoginWritesTheTokenFile(t *testing.T) {
 }
 
 func TestLoginFailsWhenTheCodeIsDenied(t *testing.T) {
-	server := newFakeAuth(t)
+	server := newChatGPTFake(t)
 	server.deny = true
-	auth := newTestAuth(t, server)
+	auth := newTestChatGPTAuth(t, server)
 
 	err := auth.Login(t.Context(), io.Discard)
 	require.ErrorContains(t, err, "HTTP 400")
@@ -64,15 +64,15 @@ func TestLoginFailsWhenTheCodeIsDenied(t *testing.T) {
 }
 
 func TestAuthorizeUsesAFreshTokenWithoutRefresh(t *testing.T) {
-	server := newFakeAuth(t)
-	auth := newTestAuth(t, server)
-	writeTestTokens(t, auth, server.token(time.Hour), "refresh-0")
+	server := newChatGPTFake(t)
+	auth := newTestChatGPTAuth(t, server)
+	chatgptWriteTestTokens(t, auth, server.token(time.Hour), "refresh-0")
 
 	header := http.Header{}
 	err := auth.Authorize(t.Context(), header)
 	require.NoError(t, err)
 	require.Equal(t, "Bearer "+server.access[0], header.Get("Authorization"))
-	require.Equal(t, testAccount, header.Get("ChatGPT-Account-Id"))
+	require.Equal(t, chatgptTestAccount, header.Get("ChatGPT-Account-Id"))
 	require.Equal(t, 0, server.refreshes)
 }
 
@@ -96,10 +96,10 @@ func TestAuthorizeRefreshesNearTheExpiry(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
-			server := newFakeAuth(t)
+			server := newChatGPTFake(t)
 			server.newRefresh = tc.newRefresh
-			auth := newTestAuth(t, server)
-			writeTestTokens(t, auth, server.token(59*time.Second), "refresh-0")
+			auth := newTestChatGPTAuth(t, server)
+			chatgptWriteTestTokens(t, auth, server.token(59*time.Second), "refresh-0")
 
 			header := http.Header{}
 			err := auth.Authorize(t.Context(), header)
@@ -107,10 +107,10 @@ func TestAuthorizeRefreshesNearTheExpiry(t *testing.T) {
 			require.Equal(t, 1, server.refreshes)
 			require.Equal(t, "refresh-0", server.refreshed)
 
-			tokens := readTokenFile(t, auth.path)
+			tokens := chatgptReadTokenFile(t, auth.path)
 			require.Equal(t, server.access[1], tokens.AccessToken)
 			require.Equal(t, tc.wantRefresh, tokens.RefreshToken)
-			require.Equal(t, testAccount, tokens.AccountID)
+			require.Equal(t, chatgptTestAccount, tokens.AccountID)
 			require.Equal(t, "Bearer "+server.access[1], header.Get("Authorization"))
 
 			err = auth.Authorize(t.Context(), header)
@@ -121,10 +121,10 @@ func TestAuthorizeRefreshesNearTheExpiry(t *testing.T) {
 }
 
 func TestAuthorizeRunsOneRefreshAtATime(t *testing.T) {
-	server := newFakeAuth(t)
+	server := newChatGPTFake(t)
 	server.newRefresh = "refresh-new"
-	auth := newTestAuth(t, server)
-	writeTestTokens(t, auth, server.token(time.Second), "refresh-0")
+	auth := newTestChatGPTAuth(t, server)
+	chatgptWriteTestTokens(t, auth, server.token(time.Second), "refresh-0")
 
 	var group sync.WaitGroup
 	errs := make(chan error, 8)
@@ -144,19 +144,19 @@ func TestAuthorizeRunsOneRefreshAtATime(t *testing.T) {
 }
 
 func TestAuthorizeNeedsALogin(t *testing.T) {
-	server := newFakeAuth(t)
-	auth := newTestAuth(t, server)
+	server := newChatGPTFake(t)
+	auth := newTestChatGPTAuth(t, server)
 
 	err := auth.Authorize(t.Context(), http.Header{})
-	require.ErrorIs(t, err, errLoginNeeded)
+	require.ErrorIs(t, err, errChatGPTLoginNeeded)
 
 	server.failRefresh = true
-	writeTestTokens(t, auth, server.token(-time.Minute), "refresh-0")
+	chatgptWriteTestTokens(t, auth, server.token(-time.Minute), "refresh-0")
 	before, err := os.ReadFile(auth.path)
 	require.NoError(t, err)
 
 	err = auth.Authorize(t.Context(), http.Header{})
-	require.ErrorIs(t, err, errLoginNeeded)
+	require.ErrorIs(t, err, errChatGPTLoginNeeded)
 	require.NotContains(t, err.Error(), "refresh-0")
 
 	after, err := os.ReadFile(auth.path)
@@ -165,9 +165,9 @@ func TestAuthorizeNeedsALogin(t *testing.T) {
 }
 
 func TestForcedRefreshLetsTheNextCallWork(t *testing.T) {
-	server := newFakeAuth(t)
-	auth := newTestAuth(t, server)
-	writeTestTokens(t, auth, server.token(-time.Hour), "refresh-0")
+	server := newChatGPTFake(t)
+	auth := newTestChatGPTAuth(t, server)
+	chatgptWriteTestTokens(t, auth, server.token(-time.Hour), "refresh-0")
 
 	responses := newModelFake(t)
 	client := newModelClient(responses.URL, auth, "agent-key", filepath.Join(t.TempDir(), "usage.jsonl"))
@@ -177,8 +177,8 @@ func TestForcedRefreshLetsTheNextCallWork(t *testing.T) {
 	require.Equal(t, 1, server.refreshes)
 	call := responses.Call(0)
 	require.Equal(t, "Bearer "+server.access[1], call.header.Get("Authorization"))
-	require.Equal(t, testAccount, call.header.Get("ChatGPT-Account-Id"))
-	require.Equal(t, server.access[1], readTokenFile(t, auth.path).AccessToken)
+	require.Equal(t, chatgptTestAccount, call.header.Get("ChatGPT-Account-Id"))
+	require.Equal(t, server.access[1], chatgptReadTokenFile(t, auth.path).AccessToken)
 }
 
 func TestModelEndpointFollowsTheConfig(t *testing.T) {
@@ -189,8 +189,8 @@ func TestModelEndpointFollowsTheConfig(t *testing.T) {
 	address, auth, err := modelEndpoint(config, dir)
 	require.NoError(t, err)
 	require.Equal(t, modelChatGPTURL, address)
-	require.IsType(t, &chatgptAuth{}, auth)
-	require.Equal(t, filepath.Join(dir, tokenFileName), auth.(*chatgptAuth).path)
+	require.IsType(t, &ChatGPTAuth{}, auth)
+	require.Equal(t, filepath.Join(dir, chatgptTokenFileName), auth.(*ChatGPTAuth).path)
 
 	config.Endpoint = endpointAPI
 	address, auth, err = modelEndpoint(config, dir)
@@ -203,10 +203,10 @@ func TestModelEndpointFollowsTheConfig(t *testing.T) {
 	require.ErrorContains(t, err, "OPENAI_API_KEY")
 }
 
-// fakeAuth is an httptest login service. The device token endpoint answers
+// 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 fakeAuth struct {
+type ChatGPTFake struct {
 	*httptest.Server
 	mutex       sync.Mutex
 	pending     int
@@ -220,36 +220,52 @@ type fakeAuth struct {
 	access      []string
 }
 
-func newFakeAuth(t *testing.T) *fakeAuth {
-	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()
 
-	f := &fakeAuth{}
-	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.newToken(ttl)
+}
 
-	return f
+// 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 *fakeAuth) userCode(w http.ResponseWriter, r *http.Request) {
+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"] != codexClientID {
+	if err != nil || body["client_id"] != chatgptClientID {
 		http.Error(w, "bad client", http.StatusBadRequest)
 		return
 	}
 
-	writeJSON(w, map[string]any{
+	chatgptWriteJSON(w, map[string]any{
 		"device_auth_id": "device-1",
 		"user_code":      "CODE-1234",
 		"interval":       "0",
 	})
 }
 
-func (f *fakeAuth) deviceToken(w http.ResponseWriter, r *http.Request) {
+func (f *ChatGPTFake) deviceToken(w http.ResponseWriter, r *http.Request) {
 	f.mutex.Lock()
 	defer f.mutex.Unlock()
 
@@ -271,14 +287,14 @@ func (f *fakeAuth) deviceToken(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	writeJSON(w, map[string]any{
+	chatgptWriteJSON(w, map[string]any{
 		"authorization_code": "auth-code",
 		"code_challenge":     "challenge",
 		"code_verifier":      "verifier",
 	})
 }
 
-func (f *fakeAuth) oauthToken(w http.ResponseWriter, r *http.Request) {
+func (f *ChatGPTFake) oauthToken(w http.ResponseWriter, r *http.Request) {
 	f.mutex.Lock()
 	defer f.mutex.Unlock()
 
@@ -290,7 +306,7 @@ func (f *fakeAuth) oauthToken(w http.ResponseWriter, r *http.Request) {
 		}
 
 		f.exchange = r.PostForm
-		writeJSON(w, map[string]any{
+		chatgptWriteJSON(w, map[string]any{
 			"access_token":  f.newToken(time.Hour),
 			"refresh_token": "refresh-1",
 			"id_token":      f.newToken(time.Hour),
@@ -301,7 +317,7 @@ func (f *fakeAuth) 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"] != codexClientID {
+	if err != nil || body["grant_type"] != "refresh_token" || body["client_id"] != chatgptClientID {
 		http.Error(w, "bad refresh", http.StatusBadRequest)
 		return
 	}
@@ -322,70 +338,54 @@ func (f *fakeAuth) oauthToken(w http.ResponseWriter, r *http.Request) {
 		answer["refresh_token"] = f.newRefresh
 	}
 
-	writeJSON(w, answer)
+	chatgptWriteJSON(w, answer)
 }
 
-// token returns a new access token that expires after ttl, for a token file
-// that a test writes.
-func (f *fakeAuth) 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)
+	f.Server = httptest.NewServer(mux)
+	t.Cleanup(f.Close)
 
-// 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 *fakeAuth) newToken(ttl time.Duration) string {
-	claims, err := json.Marshal(map[string]any{
-		"exp": time.Now().Add(ttl).Unix(),
-		authClaim: map[string]any{
-			"chatgpt_account_id": testAccount,
-		},
-		"serial": len(f.access),
-	})
-	if err != nil {
-		panic(err)
-	}
+	return f
+}
 
-	encode := base64.RawURLEncoding.EncodeToString
-	token := encode([]byte(`{"alg":"none"}`)) + "." + encode(claims) + "." + encode([]byte("fake"))
-	f.access = append(f.access, token)
+func newTestChatGPTAuth(t *testing.T, server *ChatGPTFake) *ChatGPTAuth {
+	t.Helper()
 
-	return token
+	return newChatGPTAuth(filepath.Join(t.TempDir(), chatgptTokenFileName), server.URL)
 }
 
-func newTestAuth(t *testing.T, server *fakeAuth) *chatgptAuth {
+func chatgptReadTokenFile(t *testing.T, path string) ChatGPTTokenFile {
 	t.Helper()
 
-	return newChatGPTAuth(filepath.Join(t.TempDir(), tokenFileName), server.URL)
+	data, err := os.ReadFile(path)
+	require.NoError(t, err)
+
+	var tokens ChatGPTTokenFile
+	require.NoError(t, json.Unmarshal(data, &tokens))
+
+	return tokens
 }
 
-func writeTestTokens(t *testing.T, auth *chatgptAuth, access, refresh string) {
+func chatgptWriteTestTokens(t *testing.T, auth *ChatGPTAuth, access, refresh string) {
 	t.Helper()
 
-	err := auth.write(tokenFile{
+	err := auth.write(ChatGPTTokenFile{
 		AccessToken:  access,
 		RefreshToken: refresh,
 		IDToken:      access,
-		AccountID:    testAccount,
+		AccountID:    chatgptTestAccount,
 	})
 	require.NoError(t, err)
 }
 
-func readTokenFile(t *testing.T, path string) tokenFile {
-	t.Helper()
-
-	data, err := os.ReadFile(path)
-	require.NoError(t, err)
-
-	var tokens tokenFile
-	require.NoError(t, json.Unmarshal(data, &tokens))
-
-	return tokens
-}
-
-func writeJSON(w http.ResponseWriter, value any) {
+func chatgptWriteJSON(w http.ResponseWriter, value any) {
 	w.Header().Set("Content-Type", "application/json")
 	err := json.NewEncoder(w).Encode(value)
 	if err != nil {