From 44f9b590ebe61c280cc4522e9ad7a30d550d10aa Mon Sep 17 00:00:00 2001
From: Akhilesh Kumar <akhilbussiness@gmail.com>
Date: Wed, 15 Apr 2026 17:26:18 +0000
Subject: [PATCH] test: add integration test for Gemma 4 routing

---
 integration-tests/gemma4-routing.responses |  1 +
 integration-tests/gemma4-routing.test.ts   | 52 ++++++++++++++++++++++
 2 files changed, 53 insertions(+)
 create mode 100644 integration-tests/gemma4-routing.responses
 create mode 100644 integration-tests/gemma4-routing.test.ts

diff --git a/integration-tests/gemma4-routing.responses b/integration-tests/gemma4-routing.responses
new file mode 100644
index 0000000000..3671dd4368
--- /dev/null
+++ b/integration-tests/gemma4-routing.responses
@@ -0,0 +1 @@
+{"method":"generateContentStream","response":[{"candidates":[{"content":{"parts":[{"text":"I am Gemma 4."}],"role":"model"},"index":0,"finishReason":"STOP"}]}]}
diff --git a/integration-tests/gemma4-routing.test.ts b/integration-tests/gemma4-routing.test.ts
new file mode 100644
index 0000000000..67f4400616
--- /dev/null
+++ b/integration-tests/gemma4-routing.test.ts
@@ -0,0 +1,52 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { describe, it, expect, beforeEach, afterEach } from 'vitest';
+import { join, dirname } from 'node:path';
+import { fileURLToPath } from 'node:url';
+import { TestRig } from './test-helper.js';
+
+describe('gemma4 routing', () => {
+  let rig: TestRig;
+
+  beforeEach(() => {
+    rig = new TestRig();
+  });
+
+  afterEach(async () => {
+    await rig.cleanup();
+  });
+
+  it('routes gemini-pro to gemma-4-31b-it when configured', async () => {
+    await rig.setup('gemma4-routing', {
+      fakeResponsesPath: join(
+        dirname(fileURLToPath(import.meta.url)),
+        'gemma4-routing.responses',
+      ),
+      settings: {
+        model: {
+          gemma4Variant: 'gemma-4-31b-it',
+        },
+      },
+    });
+
+    // We don't need real responses since we're just checking the routing/telemetry
+    // But TestRig might require them if it actually tries to call the API.
+    // Let's use a simple prompt.
+    await rig.run({
+      args: ['--model', 'gemini-2.5-pro', 'Hello!'],
+    });
+
+    const hasApiRequestEvent = await rig.waitForTelemetryEvent('api_request');
+    expect(hasApiRequestEvent).toBe(true);
+
+    const lastRequest = rig.readLastApiRequest();
+
+    expect(lastRequest).not.toBeNull();
+    // The telemetry logger records the final requested model ID in the 'model' attribute
+    expect(lastRequest?.attributes?.model).toBe('gemma-4-31b-it');
+  });
+});