Re-add ToolTesterMixin

huggingface · Sep 13, 2024 · d6f703f · d6f703f
1 parent 98f5463
commit d6f703f
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/tests/agents/test_final_answer.py b/tests/agents/test_final_answer.py
@@ -24,12 +24,14 @@
 from transformers.agents.default_tools import FinalAnswerTool
 from transformers.testing_utils import get_tests_dir, require_torch
 
+from .test_tools_common import ToolTesterMixin
+
 
 if is_torch_available():
     import torch
 
 
-class FinalAnswerToolTester(unittest.TestCase):
+class FinalAnswerToolTester(unittest.TestCase, ToolTesterMixin):
     def setUp(self):
         self.inputs = {"answer": "Final answer"}
         self.tool = FinalAnswerTool()