From 2b1a791a0b24be48902f0cbbf1c617ee53b69bd1 Mon Sep 17 00:00:00 2001
From: christine betts <chrstn@uw.edu>
Date: Tue, 2 Dec 2025 15:35:58 -0500
Subject: [PATCH] Use polling for extensions-reload integration test (#14391)

---
 integration-tests/extensions-reload.test.ts | 56 +++++++++++++++------
 integration-tests/test-helper.ts            | 27 ++++++++++
 2 files changed, 68 insertions(+), 15 deletions(-)

diff --git a/integration-tests/extensions-reload.test.ts b/integration-tests/extensions-reload.test.ts
index 91e0d9c0ad..43d3b4a278 100644
--- a/integration-tests/extensions-reload.test.ts
+++ b/integration-tests/extensions-reload.test.ts
@@ -13,6 +13,8 @@ import { safeJsonStringify } from '@google/gemini-cli-core/src/utils/safeJsonStr
 import { env } from 'node:process';
 import { platform } from 'node:os';
 
+import stripAnsi from 'strip-ansi';
+
 const itIf = (condition: boolean) => (condition ? it : it.skip);
 
 describe('extension reloading', () => {
@@ -76,13 +78,22 @@ describe('extension reloading', () => {
       await run.expectText(
         'test-extension (v0.0.1) - active (update available)',
       );
-      // Wait a bit for the UI to settle
-      await new Promise((resolve) => setTimeout(resolve, 3000));
-      await run.sendKeys('\u0015/mcp list\r');
-      await run.expectText(
-        'test-server (from test-extension) - Ready (1 tool)',
+      // Wait for the UI to settle and retry the command until we see the update
+      await new Promise((resolve) => setTimeout(resolve, 1000));
+
+      // Poll for the updated list
+      await rig.pollCommand(
+        () => run.sendKeys('\u0015/mcp list\r'),
+        () => {
+          const output = stripAnsi(run.output);
+          return (
+            output.includes(
+              'test-server (from test-extension) - Ready (1 tool)',
+            ) && output.includes('- hello')
+          );
+        },
+        30000, // 30s timeout
       );
-      await run.expectText('- hello');
 
       // Update the extension, expect the list to update, and mcp servers as well.
       await run.sendKeys('\u0015/extensions update test-extension');
@@ -97,16 +108,31 @@ describe('extension reloading', () => {
       await run.expectText(
         'Extension "test-extension" successfully updated: 0.0.1 → 0.0.2',
       );
-      await new Promise((resolve) => setTimeout(resolve, 3000));
-      await run.sendKeys('\u0015/extensions list\r');
-      await run.expectText('test-extension (v0.0.2) - active (updated)');
-      // Wait a bit for the UI to settle
-      await new Promise((resolve) => setTimeout(resolve, 3000));
-      await run.sendKeys('\u0015/mcp list\r');
-      await run.expectText(
-        'test-server (from test-extension) - Ready (1 tool)',
+
+      // Poll for the updated extension version
+      await rig.pollCommand(
+        () => run.sendKeys('\u0015/extensions list\r'),
+        () =>
+          stripAnsi(run.output).includes(
+            'test-extension (v0.0.2) - active (updated)',
+          ),
+        30000,
       );
-      await run.expectText('- goodbye');
+
+      // Poll for the updated mcp tool
+      await rig.pollCommand(
+        () => run.sendKeys('\u0015/mcp list\r'),
+        () => {
+          const output = stripAnsi(run.output);
+          return (
+            output.includes(
+              'test-server (from test-extension) - Ready (1 tool)',
+            ) && output.includes('- goodbye')
+          );
+        },
+        30000,
+      );
+
       await run.sendText('/quit');
       await run.sendKeys('\r');
 
diff --git a/integration-tests/test-helper.ts b/integration-tests/test-helper.ts
index 418dd57ef0..246af4bb56 100644
--- a/integration-tests/test-helper.ts
+++ b/integration-tests/test-helper.ts
@@ -159,6 +159,14 @@ interface ParsedLog {
     success?: boolean;
     duration_ms?: number;
     request_text?: string;
+    hook_event_name?: string;
+    hook_name?: string;
+    hook_input?: Record<string, unknown>;
+    hook_output?: Record<string, unknown>;
+    exit_code?: number;
+    stdout?: string;
+    stderr?: string;
+    error?: string;
   };
   scopeMetrics?: {
     metrics: {
@@ -1081,4 +1089,23 @@ export class TestRig {
 
     return logs;
   }
+
+  async pollCommand(
+    commandFn: () => Promise<void>,
+    predicateFn: () => boolean,
+    timeout: number = 30000,
+    interval: number = 1000,
+  ) {
+    const startTime = Date.now();
+    while (Date.now() - startTime < timeout) {
+      await commandFn();
+      // Give it a moment to process
+      await new Promise((resolve) => setTimeout(resolve, 500));
+      if (predicateFn()) {
+        return;
+      }
+      await new Promise((resolve) => setTimeout(resolve, interval));
+    }
+    throw new Error(`pollCommand timed out after ${timeout}ms`);
+  }
 }