From 1c2f0f708c4e30ed706abcf8c110eabdd25720c1 Mon Sep 17 00:00:00 2001
From: Xuan Son Nguyen <son@huggingface.co>
Date: Wed, 20 Nov 2024 19:24:24 +0100
Subject: [PATCH] fix save slot test

---
 examples/server/tests/unit/test_slot_save.py | 7 ++++---
 examples/server/tests/utils.py               | 2 +-
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/examples/server/tests/unit/test_slot_save.py b/examples/server/tests/unit/test_slot_save.py
index 68249b0be..38704f5ec 100644
--- a/examples/server/tests/unit/test_slot_save.py
+++ b/examples/server/tests/unit/test_slot_save.py
@@ -8,6 +8,7 @@ def create_server():
     global server
     server = ServerPreset.tinyllama2()
     server.slot_save_path = "./tmp"
+    server.temperature = 0.0
 
 
 def test_slot_save_restore():
@@ -21,7 +22,7 @@ def test_slot_save_restore():
         "cache_prompt": True,
     })
     assert res.status_code == 200
-    assert match_regex("(Lily|cake)+", res.body["content"])
+    assert match_regex("(Whiskers|Flana)+", res.body["content"])
     assert res.body["timings"]["prompt_n"] == 21  # all tokens are processed
 
     # Save state of slot 1
@@ -79,7 +80,7 @@ def test_slot_erase():
         "cache_prompt": True,
     })
     assert res.status_code == 200
-    assert match_regex("(Lily|cake)+", res.body["content"])
+    assert match_regex("(Whiskers|Flana)+", res.body["content"])
     assert res.body["timings"]["prompt_n"] == 21  # all tokens are processed
 
     # erase slot 1
@@ -93,5 +94,5 @@ def test_slot_erase():
         "cache_prompt": True,
     })
     assert res.status_code == 200
-    assert match_regex("(Lily|cake)+", res.body["content"])
+    assert match_regex("(Whiskers|Flana)+", res.body["content"])
     assert res.body["timings"]["prompt_n"] == 21  # all tokens are processed
diff --git a/examples/server/tests/utils.py b/examples/server/tests/utils.py
index 03e0a4a97..0808a92b2 100644
--- a/examples/server/tests/utils.py
+++ b/examples/server/tests/utils.py
@@ -47,7 +47,7 @@ class ServerProcess:
     model_url: str | None = None
     model_file: str | None = None
     n_threads: int | None = None
-    n_gpu_layer: str | None = None
+    n_gpu_layer: int | None = None
     n_batch: int | None = None
     n_ubatch: int | None = None
     n_ctx: int | None = None