Skip to content

Commit 5bf29af

Browse files
ggerganovngxson
andcommitted
tests : improve "tokens" type check
Co-authored-by: Xuan Son Nguyen <[email protected]>
1 parent 8bcfc55 commit 5bf29af

File tree

1 file changed

+4
-2
lines changed

1 file changed

+4
-2
lines changed

examples/server/tests/unit/test_completion.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,8 @@ def test_completion(prompt: str, n_predict: int, re_content: str, n_prompt: int,
2929
assert type(res.body["has_new_line"]) == bool
3030
assert match_regex(re_content, res.body["content"])
3131
if return_tokens:
32-
assert res.body["tokens"] != []
32+
assert len(res.body["tokens"]) > 0
33+
assert all(type(tok) == int for tok in res.body["tokens"])
3334
else:
3435
assert res.body["tokens"] == []
3536

@@ -61,7 +62,8 @@ def test_completion_stream(prompt: str, n_predict: int, re_content: str, n_promp
6162
assert data["generation_settings"]["seed"] == server.seed
6263
assert match_regex(re_content, content)
6364
else:
64-
assert data["tokens"] != []
65+
assert len(res.body["tokens"]) > 0
66+
assert all(type(tok) == int for tok in res.body["tokens"])
6567
content += data["content"]
6668

6769

0 commit comments

Comments
 (0)