Update tests

neuml · Nov 5, 2024 · 058ca9b · 058ca9b
1 parent 9ddfc41
commit 058ca9b
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 10 deletions.
diff --git a/test/python/testagent.py b/test/python/testagent.py
@@ -49,20 +49,25 @@ def testToolsBasic(self):
         Test adding basic function tools
         """
 
-        def today(iso):
+        class DateTime:
+            """
+            Date time instance
             """
-            Gets the current date and time
 
-            Args:
-                iso: date will be converted to iso format if True
+            def __call__(self, iso):
+                """
+                Gets the current date and time
 
-            Returns:
-                current date and time
-            """
+                Args:
+                    iso: date will be converted to iso format if True
+
+                Returns:
+                    current date and time
+                """
 
-            return datetime.today().isoformat() if iso else datetime.today()
+                return datetime.today().isoformat() if iso else datetime.today()
 
-        today = {"name": "today", "description": "Gets the current date and time", "target": today}
+        today = {"name": "today", "description": "Gets the current date and time", "target": DateTime()}
 
         def current() -> str:
             """

diff --git a/test/python/testpipeline/testllm/testllm.py b/test/python/testpipeline/testllm/testllm.py
@@ -74,7 +74,7 @@ def testMaxLength(self):
         """
 
         model = LLM("sshleifer/tiny-gpt2")
-        self.assertIsNotNone(model("Hello, how are", maxlength=10))
+        self.assertIsInstance(model("Hello, how are", maxlength=10), str)
 
     def testNotImplemented(self):
         """