ensure output format on chat enabled model tests

Signed-off-by: Jeffrey Martin <[email protected]>
NVIDIA · Dec 19, 2024 · 86de116 · 86de116
1 parent 900fc8f
commit 86de116
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/tests/generators/test_huggingface.py b/tests/generators/test_huggingface.py
@@ -58,8 +58,11 @@ def test_pipeline_chat(mocker, hf_generator_config):
     mock_format = mocker.patch.object(
         g, "_format_chat_prompt", wraps=g._format_chat_prompt
     )
-    g.generate("Hello world!")
+    output = g.generate("Hello world!")
     mock_format.assert_called_once()
+    assert len(output) == 1
+    for item in output:
+        assert isinstance(item, str)
 
 
 def test_inference(mocker, hf_mock_response, hf_generator_config):
@@ -141,8 +144,11 @@ def test_model_chat(mocker, hf_generator_config):
     mock_format = mocker.patch.object(
         g, "_format_chat_prompt", wraps=g._format_chat_prompt
     )
-    g.generate("Hello world!")
+    output = g.generate("Hello world!")
     mock_format.assert_called_once()
+    assert len(output) == 1
+    for item in output:
+        assert isinstance(item, str)
 
 
 def test_select_hf_device():