From e09a78a25dec3f42175ad59ecc540a54bed0b9b2 Mon Sep 17 00:00:00 2001
From: Pascal <admin@serveurperso.com>
Date: Thu, 29 Oct 2026 01:35:01 +0100
Subject: [PATCH] tests: fix CI flakes in server cancel and sleep tests
 (#31330)

test_cancel_request enables context shift so the generation only ends
when the request is cancelled, then polls /slots with a deadline.
Without context shift the generation fills the context and frees the
slot on its own on fast machines, so the test passes even with
cancellation disabled.

test_server_sleep uses the existing wait_for_sleep helper instead of a
fixed sleep.
---
 tools/server/tests/unit/test_completion.py | 15 ++++++++++-----
 tools/server/tests/unit/test_sleep.py      |  3 +--
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/tools/server/tests/unit/test_completion.py b/tools/server/tests/unit/test_completion.py
index 572f5368c..6afcb1773 100644
--- a/tools/server/tests/unit/test_completion.py
+++ b/tools/server/tests/unit/test_completion.py
@@ -637,18 +637,23 @@ def test_cancel_request():
     server.n_predict = -1
     server.n_slots = 1
     server.server_slots = True
+    server.enable_ctx_shift = True
     server.start()
-    # send a request that will take a long time, but cancel it before it finishes
+    # context shift keeps the generation running until the request is cancelled
     try:
         server.make_request("POST", "/completion", data={
             "prompt": "I believe the meaning of life is",
         }, timeout=0.1)
     except requests.exceptions.ReadTimeout:
         pass # expected
-    # make sure the slot is free
-    time.sleep(2)
-    res = server.make_request("GET", "/slots")
-    assert res.body[0]["is_processing"] == False
+    # the slot is freed once the server notices the closed connection
+    deadline = time.time() + 10
+    while time.time() < deadline:
+        res = server.make_request("GET", "/slots")
+        if not res.body[0]["is_processing"]:
+            return
+        time.sleep(0.1)
+    raise AssertionError("slot is still processing after the request was cancelled")
 
 
 # this test exercises the host-memory prompt cache
diff --git a/tools/server/tests/unit/test_sleep.py b/tools/server/tests/unit/test_sleep.py
index 4edb26450..b3ac937f6 100644
--- a/tools/server/tests/unit/test_sleep.py
+++ b/tools/server/tests/unit/test_sleep.py
@@ -45,8 +45,7 @@ def test_server_sleep():
     server.sleep_idle_seconds = 1
     server.start()
 
-    # wait a bit so that server can go to sleep
-    time.sleep(2)
+    wait_for_sleep(server)
 
     # make sure these endpoints are still responsive after sleep
     res = server.make_request("GET", "/health")