From 5bf29af841ed7e58077aff136c1a6d3ba59af115 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Wed, 18 Dec 2024 10:02:01 +0200 Subject: [PATCH] tests : improve "tokens" type check Co-authored-by: Xuan Son Nguyen --- examples/server/tests/unit/test_completion.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/examples/server/tests/unit/test_completion.py b/examples/server/tests/unit/test_completion.py index 39b5c8dbe30cf..135cd6d0c932b 100644 --- a/examples/server/tests/unit/test_completion.py +++ b/examples/server/tests/unit/test_completion.py @@ -29,7 +29,8 @@ def test_completion(prompt: str, n_predict: int, re_content: str, n_prompt: int, assert type(res.body["has_new_line"]) == bool assert match_regex(re_content, res.body["content"]) if return_tokens: - assert res.body["tokens"] != [] + assert len(res.body["tokens"]) > 0 + assert all(type(tok) == int for tok in res.body["tokens"]) else: assert res.body["tokens"] == [] @@ -61,7 +62,8 @@ def test_completion_stream(prompt: str, n_predict: int, re_content: str, n_promp assert data["generation_settings"]["seed"] == server.seed assert match_regex(re_content, content) else: - assert data["tokens"] != [] + assert len(res.body["tokens"]) > 0 + assert all(type(tok) == int for tok in res.body["tokens"]) content += data["content"]