test(grok): cover OAuth chat permission identity

2026-09-24 16:05:44 +08:00 · 2026-07-13 10:11:33 +08:00
parent ce3f12bbff
commit c4ff604e93
1 changed files with 47 additions and 38 deletions
@@ -15,49 +15,58 @@ import (
 	"github.com/stretchr/testify/suite"
 )

-func TestHTTPUpstreamDoAppliesGrokCLIIdentityBeforeRoundTrip(t *testing.T) {
+func TestHTTPUpstreamDoAppliesGrokCLIIdentityBeforeOAuthRoundTrip(t *testing.T) {
 	t.Setenv("XAI_GROK_CLI_VERSION", "")

-	upstream := NewHTTPUpstream(nil)
-	svc, ok := upstream.(*httpUpstreamService)
-	require.True(t, ok)
+	for _, endpoint := range []string{"responses", "chat/completions"} {
+		t.Run(endpoint, func(t *testing.T) {
+			upstream := NewHTTPUpstream(nil)
+			svc, ok := upstream.(*httpUpstreamService)
+			require.True(t, ok)

-	const accountID int64 = 4079
-	isolation := svc.getIsolationMode()
-	profile := service.HTTPUpstreamProfileDefault
-	proxyKey := directProxyKey
-	protocolMode := svc.resolveProtocolMode(profile, proxyKey, nil)
-	settings := svc.resolvePoolSettings(isolation, 1)
-	settings = svc.applyProfilePoolSettings(settings, profile)
-	cacheKey := buildCacheKey(isolation, proxyKey, accountID, protocolMode)
+			const accountID int64 = 4084
+			isolation := svc.getIsolationMode()
+			profile := service.HTTPUpstreamProfileDefault
+			proxyKey := directProxyKey
+			protocolMode := svc.resolveProtocolMode(profile, proxyKey, nil)
+			settings := svc.resolvePoolSettings(isolation, 1)
+			settings = svc.applyProfilePoolSettings(settings, profile)
+			cacheKey := buildCacheKey(isolation, proxyKey, accountID, protocolMode)

-	var capturedHeaders http.Header
-	svc.clients[cacheKey] = &upstreamClientEntry{
-		client: &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
-			capturedHeaders = req.Header.Clone()
-			return &http.Response{
-				StatusCode: http.StatusOK,
-				Header:     make(http.Header),
-				Body:       http.NoBody,
-				Request:    req,
-			}, nil
-		})},
-		proxyKey:     proxyKey,
-		poolKey:      buildPoolKey(settings, protocolMode),
-		protocolMode: protocolMode,
+			var capturedHeaders http.Header
+			svc.clients[cacheKey] = &upstreamClientEntry{
+				client: &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
+					capturedHeaders = req.Header.Clone()
+					statusCode := http.StatusOK
+					if req.Header.Get("X-XAI-Token-Auth") != "xai-grok-cli" {
+						statusCode = http.StatusForbidden
+					}
+					return &http.Response{
+						StatusCode: statusCode,
+						Header:     make(http.Header),
+						Body:       http.NoBody,
+						Request:    req,
+					}, nil
+				})},
+				proxyKey:     proxyKey,
+				poolKey:      buildPoolKey(settings, protocolMode),
+				protocolMode: protocolMode,
+			}
+
+			req, err := http.NewRequest(http.MethodPost, "https://cli-chat-proxy.grok.com/v1/"+endpoint, nil)
+			require.NoError(t, err)
+			req.Header.Set("User-Agent", "sub2api-grok/1.0")
+
+			resp, err := svc.Do(req, "", accountID, 1)
+			require.NoError(t, err)
+			require.Equal(t, http.StatusOK, resp.StatusCode)
+			require.NoError(t, resp.Body.Close())
+
+			require.Equal(t, "0.2.93", capturedHeaders.Get("x-grok-client-version"))
+			require.Equal(t, "xai-grok-cli", capturedHeaders.Get("X-XAI-Token-Auth"))
+			require.Equal(t, "xai-grok-workspace/0.2.93", capturedHeaders.Get("User-Agent"))
+		})
 	}
-
-	req, err := http.NewRequest(http.MethodPost, "https://cli-chat-proxy.grok.com/v1/responses", nil)
-	require.NoError(t, err)
-	req.Header.Set("User-Agent", "sub2api-grok/1.0")
-
-	resp, err := svc.Do(req, "", accountID, 1)
-	require.NoError(t, err)
-	require.NoError(t, resp.Body.Close())
-
-	require.Equal(t, "0.2.93", capturedHeaders.Get("x-grok-client-version"))
-	require.Equal(t, "xai-grok-cli", capturedHeaders.Get("X-XAI-Token-Auth"))
-	require.Equal(t, "xai-grok-workspace/0.2.93", capturedHeaders.Get("User-Agent"))
 }

 func TestApplyGrokCLIProxyHeaders(t *testing.T) {