refactor: better test for answer

umbertogriffo · May 4, 2024 · dd42c4e · dd42c4e
1 parent 9410c91
commit dd42c4e
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/tests/test_ctransformer_client.py b/tests/test_ctransformer_client.py
@@ -57,6 +57,6 @@ def test_encode_prompt(ctransformers_client):
 
 
 def test_generate_answer(ctransformers_client):
-    prompt = "Tell me a joke"
+    prompt = "What is the capital city of Italy?"
     generated_answer = ctransformers_client.generate_answer(prompt, max_new_tokens=10)
-    assert generated_answer is not None
+    assert "rome" in generated_answer.lower()
diff --git a/tests/test_lamacpp_client.py b/tests/test_lamacpp_client.py
@@ -38,6 +38,6 @@ def test_init_raises_value_error_for_invalid_client_type(mock_model_folder, inva
 
 
 def test_generate_answer(lamacpp_client):
-    prompt = "Tell me a joke"
+    prompt = "What is the capital city of Italy?"
     generated_answer = lamacpp_client.generate_answer(prompt, max_new_tokens=10)
-    assert generated_answer is not None
+    assert "rome" in generated_answer.lower()