From dd42c4e4985ede2fe1c8cb5438fe1b1dcba55adf Mon Sep 17 00:00:00 2001
From: Umberto Griffo <1609440+umbertogriffo@users.noreply.github.com>
Date: Sat, 4 May 2024 13:46:17 +0100
Subject: [PATCH] refactor: better test for answer

---
 tests/test_ctransformer_client.py | 4 ++--
 tests/test_lamacpp_client.py      | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/tests/test_ctransformer_client.py b/tests/test_ctransformer_client.py
index 928695e..6c33d72 100644
--- 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
index 82fc044..70947ff 100644
--- 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()