From 9233271823d1022bd174cf8d3fc5002e6b3aac3e Mon Sep 17 00:00:00 2001
From: Xuan Son Nguyen <son@huggingface.co>
Date: Thu, 30 Apr 2026 20:07:55 +0200
Subject: [PATCH] fix test case

---
 tools/server/tests/unit/test_compat_gcp.py | 5 +++--
 tools/server/tests/utils.py                | 8 +++++++-
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/tools/server/tests/unit/test_compat_gcp.py b/tools/server/tests/unit/test_compat_gcp.py
index 73093507ac..aba67bb353 100644
--- a/tools/server/tests/unit/test_compat_gcp.py
+++ b/tools/server/tests/unit/test_compat_gcp.py
@@ -8,9 +8,10 @@ server: ServerProcess
 def create_server():
     global server
     server = ServerPreset.tinyllama2()
+    server.gcp_compat = True
 
 
-def test_vertexai_predict_camel_case():
+def test_gcp_predict_camel_case():
     global server
     server.start()
     res = server.make_request("POST", "/predict", data={
@@ -34,7 +35,7 @@ def test_vertexai_predict_camel_case():
     assert len(prediction["choices"][0]["message"]["content"]) > 0
 
 
-def test_vertexai_predict_multiple_instances():
+def test_gcp_predict_multiple_instances():
     global server
     server.n_slots = 2
     server.start()
diff --git a/tools/server/tests/utils.py b/tools/server/tests/utils.py
index 88700487be..e21219a817 100644
--- a/tools/server/tests/utils.py
+++ b/tools/server/tests/utils.py
@@ -105,6 +105,7 @@ class ServerProcess:
     no_cache_idle_slots: bool = False
     log_path: str | None = None
     webui_mcp_proxy: bool = False
+    gcp_compat: bool = False
 
     # session variables
     process: subprocess.Popen | None = None
@@ -119,6 +120,9 @@ class ServerProcess:
         self.external_server = "DEBUG_EXTERNAL" in os.environ
 
     def start(self, timeout_seconds: int = DEFAULT_HTTP_TIMEOUT) -> None:
+        env = {**os.environ}
+        if "LLAMA_CACHE" not in os.environ:
+            env["LLAMA_CACHE"] = "tmp"
         if self.external_server:
             print(f"[external_server]: Assuming external server running on {self.server_host}:{self.server_port}")
             return
@@ -243,6 +247,8 @@ class ServerProcess:
             server_args.append("--no-cache-idle-slots")
         if self.webui_mcp_proxy:
             server_args.append("--webui-mcp-proxy")
+        if self.gcp_compat:
+            env["AIP_MODE"] = "PREDICTION"
 
         args = [str(arg) for arg in [server_path, *server_args]]
         print(f"tests: starting server with: {' '.join(args)}")
@@ -263,7 +269,7 @@ class ServerProcess:
             creationflags=flags,
             stdout=self._log,
             stderr=self._log if self._log != sys.stdout else sys.stdout,
-            env={**os.environ, "LLAMA_CACHE": "tmp"} if "LLAMA_CACHE" not in os.environ else None,
+            env=env,
         )
         server_instances.add(self)