Fix print test

argilla-io · Oct 7, 2024 · 6a79f0f · 6a79f0f
1 parent a6349da
commit 6a79f0f
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/tests/integration/test_prints.py b/tests/integration/test_prints.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from functools import partial
+from typing import Union
 
 import pytest
 
@@ -57,14 +58,14 @@
 
 
 class TestLLM(DummyLLM, MagpieChatTemplateMixin):
-    magpie_pre_query_template: str = "llama3"
+    magpie_pre_query_template: Union[str, None] = "llama3"
 
 
 llm = TestLLM()
 
 
 @pytest.mark.parametrize("task", tasks)
-def test_prints(task):
+def test_prints(task) -> None:
     t = task(llm=llm)
     t.load()
     t.print()

diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py
@@ -53,7 +53,7 @@ def model_name(self) -> str:
         return "test"
 
     def generate(  # type: ignore
-        self, input: "FormattedInput", num_generations: int = 1
+        self, inputs: "FormattedInput", num_generations: int = 1
     ) -> "GenerateOutput":
         return ["output" for _ in range(num_generations)]