From 337efa9aa3ddaaf117d498c92d5fdc19f860e872 Mon Sep 17 00:00:00 2001
From: Sehoon Shon <sshon@google.com>
Date: Mon, 23 Mar 2026 14:29:14 -0400
Subject: [PATCH] test(utils): support model fallbacks in TestRig

Support model fallbacks in integration tests by allowing a comma-separated list of models via the `GEMINI_MODEL` environment variable or the `model` option in `TestRig.setup()`. This will automatically configure a fallback chain using `dynamicModelConfiguration` with `silent` actions.

Fixes #23568
---
 packages/test-utils/src/test-rig.ts | 58 +++++++++++++++++++++++++----
 1 file changed, 51 insertions(+), 7 deletions(-)

diff --git a/packages/test-utils/src/test-rig.ts b/packages/test-utils/src/test-rig.ts
index ee091bee92..76c63428d5 100644
--- a/packages/test-utils/src/test-rig.ts
+++ b/packages/test-utils/src/test-rig.ts
@@ -348,6 +348,7 @@ export class TestRig {
   private _interactiveRuns: InteractiveRun[] = [];
   private _spawnedProcesses: ChildProcess[] = [];
   private _initialized = false;
+  private _modelOverride?: string;
 
   setup(
     testName: string,
@@ -355,9 +356,11 @@ export class TestRig {
       settings?: Record<string, unknown>;
       state?: Record<string, unknown>;
       fakeResponsesPath?: string;
+      model?: string;
     } = {},
   ) {
     this.testName = testName;
+    this._modelOverride = options.model;
     const sanitizedName = sanitizeTestName(testName);
     const testFileDir =
       env['INTEGRATION_TEST_FILE_DIR'] || join(os.tmpdir(), 'gemini-cli-tests');
@@ -429,6 +432,53 @@ export class TestRig {
     // The container mounts the test directory at the same path as the host
     const telemetryPath = join(this.homeDir!, 'telemetry.log'); // Always use home directory for telemetry
 
+    const modelEnv = this._modelOverride || env['GEMINI_MODEL'];
+    const models = modelEnv
+      ? modelEnv.split(',').map((m) => m.trim())
+      : [DEFAULT_GEMINI_MODEL];
+
+    const modelSettings: Record<string, any> = {};
+    if (models.length > 1) {
+      // Setup custom fallback chain for integration tests
+      const chainName = 'integration-test-model';
+      modelSettings['model'] = { name: chainName };
+      modelSettings['experimental'] = { dynamicModelConfiguration: true };
+      modelSettings['modelConfigs'] = {
+        modelDefinitions: {
+          [chainName]: {
+            tier: 'auto',
+            isPreview: true,
+            isVisible: false,
+          },
+        },
+        modelIdResolutions: {
+          [chainName]: {
+            default: models[0],
+          },
+        },
+        modelChains: {
+          [chainName]: models.map((m, i) => ({
+            model: m,
+            isLastResort: i === models.length - 1,
+            actions: {
+              terminal: 'silent',
+              transient: 'silent',
+              not_found: 'silent',
+              unknown: 'silent',
+            },
+            stateTransitions: {
+              terminal: 'terminal',
+              transient: 'terminal',
+              not_found: 'terminal',
+              unknown: 'terminal',
+            },
+          })),
+        },
+      };
+    } else {
+      modelSettings['model'] = { name: models[0] };
+    }
+
     const settings = deepMerge(
       {
         general: {
@@ -453,13 +503,7 @@ export class TestRig {
         ui: {
           useAlternateBuffer: true,
         },
-        ...(env['GEMINI_TEST_TYPE'] === 'integration'
-          ? {
-              model: {
-                name: DEFAULT_GEMINI_MODEL,
-              },
-            }
-          : {}),
+        ...(env['GEMINI_TEST_TYPE'] === 'integration' ? modelSettings : {}),
         sandbox:
           env['GEMINI_SANDBOX'] !== 'false' ? env['GEMINI_SANDBOX'] : false,
         // Don't show the IDE connection dialog when running from VsCode