Name and order the ChatGPT login code

[onechat] / chatgpt_test.go

diff --git a/auth_test.go b/chatgpt_test.go
similarity index 72%
rename from auth_test.go
rename to chatgpt_test.go
index 06216057cd64a3e9e8dd3ff2195c63738a3bbf93..982a6ff70e51d08164515a4f315f1c11e3ca54fd 100644 (file)
--- a/auth_test.go
+++ b/chatgpt_test.go
@@ -18,13 +18,13 @@ import (
        "github.com/stretchr/testify/require"
 )
 
        "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) {
 
 func TestLoginWritesTheTokenFile(t *testing.T) {

-       server := newFakeAuth(t)
+       server := newChatGPTFake(t)

        server.pending = 2
        server.pending = 2

-       auth := newTestAuth(t, server)
+       auth := newTestChatGPTAuth(t, server)

 
        var out strings.Builder
        err := auth.Login(t.Context(), &out)
 
        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, "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"))
 
        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, "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)
        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) {
 }
 
 func TestLoginFailsWhenTheCodeIsDenied(t *testing.T) {

-       server := newFakeAuth(t)
+       server := newChatGPTFake(t)

        server.deny = true
        server.deny = true

-       auth := newTestAuth(t, server)
+       auth := newTestChatGPTAuth(t, server)

 
        err := auth.Login(t.Context(), io.Discard)
        require.ErrorContains(t, err, "HTTP 400")
 
        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) {
 }
 
 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"))
 
        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)
 }
 
        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) {
 
        for _, tc := range testCases {
                t.Run(tc.name, func(t *testing.T) {

-                       server := newFakeAuth(t)
+                       server := newChatGPTFake(t)

                        server.newRefresh = tc.newRefresh
                        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)
 
                        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)
 
                        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, 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)
                        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) {
 }
 
 func TestAuthorizeRunsOneRefreshAtATime(t *testing.T) {

-       server := newFakeAuth(t)
+       server := newChatGPTFake(t)

        server.newRefresh = "refresh-new"
        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)
 
        var group sync.WaitGroup
        errs := make(chan error, 8)


@@ -144,19 +144,19 @@ func TestAuthorizeRunsOneRefreshAtATime(t *testing.T) {
 }
 
 func TestAuthorizeNeedsALogin(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{})
 
        err := auth.Authorize(t.Context(), http.Header{})

-       require.ErrorIs(t, err, errLoginNeeded)
+       require.ErrorIs(t, err, errChatGPTLoginNeeded)

 
        server.failRefresh = true
 
        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{})
        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)
        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) {
 }
 
 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"))
 
        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, 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) {
 }
 
 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)
        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)
 
        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")
 }
 
        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.
 // 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
        *httptest.Server
        mutex       sync.Mutex
        pending     int


@@ -220,36 +220,52 @@ type fakeAuth struct {
        access      []string
 }
 
        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)
        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
        }
 
                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",
        })
 }
 
                "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()
 
        f.mutex.Lock()
        defer f.mutex.Unlock()
 


@@ -271,14 +287,14 @@ func (f *fakeAuth) deviceToken(w http.ResponseWriter, r *http.Request) {
                return
        }
 
                return
        }
 

-       writeJSON(w, map[string]any{
+       chatgptWriteJSON(w, map[string]any{

                "authorization_code": "auth-code",
                "code_challenge":     "challenge",
                "code_verifier":      "verifier",
        })
 }
 
                "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()
 
        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
                }
 
                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),
                        "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)
 
        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
        }
                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
        }
 
                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()
 
        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()
 
        t.Helper()
 

-       err := auth.write(tokenFile{
+       err := auth.write(ChatGPTTokenFile{

                AccessToken:  access,
                RefreshToken: refresh,
                IDToken:      access,
                AccessToken:  access,
                RefreshToken: refresh,
                IDToken:      access,

-               AccountID:    testAccount,
+               AccountID:    chatgptTestAccount,

        })
        require.NoError(t, err)
 }
 
        })
        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 {
        w.Header().Set("Content-Type", "application/json")
        err := json.NewEncoder(w).Encode(value)
        if err != nil {