Refactor in preparation for Reauth (#1196)

1 files changed, 7 insertions, 2 deletions

diff --git a/packages/core/src/code_assist/oauth2.test.ts b/packages/core/src/code_assist/oauth2.test.ts
index 47bd45b3..0f5b791b 100644
--- a/packages/core/src/code_assist/oauth2.test.ts
+++ b/packages/core/src/code_assist/oauth2.test.ts
@@ -73,8 +73,10 @@ describe('oauth2', () => {
       (resolve) => (serverListeningCallback = resolve),
     );
 
+    let capturedPort = 0;
     const mockHttpServer = {
       listen: vi.fn((port: number, callback?: () => void) => {
+        capturedPort = port;
         if (callback) {
           callback();
         }
@@ -86,7 +88,7 @@ describe('oauth2', () => {
         }
       }),
       on: vi.fn(),
-      address: () => ({ port: 1234 }),
+      address: () => ({ port: capturedPort }),
     };
     vi.mocked(http.createServer).mockImplementation((cb) => {
       requestCallback = cb as http.RequestListener<
@@ -115,7 +117,10 @@ describe('oauth2', () => {
     expect(client).toBe(mockOAuth2Client);
 
     expect(open).toHaveBeenCalledWith(mockAuthUrl);
-    expect(mockGetToken).toHaveBeenCalledWith(mockCode);
+    expect(mockGetToken).toHaveBeenCalledWith({
+      code: mockCode,
+      redirect_uri: `http://localhost:${capturedPort}/oauth2callback`,
+    });
     expect(mockSetCredentials).toHaveBeenCalledWith(mockTokens);
 
     const tokenPath = path.join(tempHomeDir, '.gemini', 'oauth_creds.json');