Add tests for UDF and with session

Author: auxten

tests/test_stateful.py

@@ -93,6 +93,16 @@ def test_tmp(self):
         ret = sess2.query("SELECT chdb_xxx()", "CSV")
         self.assertEqual(str(ret), "")
 
+    def test_context_mgr(self):
+        with session.Session() as sess:
+            sess.query("CREATE FUNCTION chdb_xxx AS () -> '0.12.0'", "CSV")
+            ret = sess.query("SELECT chdb_xxx()", "CSV")
+            self.assertEqual(str(ret), '"0.12.0"\n')
+
+        with session.Session() as sess:
+            ret = sess.query("SELECT chdb_xxx()", "CSV")
+            self.assertEqual(str(ret), "")
+
     def test_zfree_thread_count(self):
         time.sleep(3)
         thread_count = current_process.num_threads()

tests/test_udf.py

@@ -1,13 +1,58 @@
 #!python3
 
-from chdb.udf import *
+import unittest
+from chdb.udf import chdb_udf
+from chdb.session import Session
+from chdb import query
 
-@to_clickhouse_udf()
+
+@chdb_udf()
 def sum_udf(lhs, rhs):
-    import time
-    time.sleep(1)
     return int(lhs) + int(rhs)
 
-@to_clickhouse_udf(return_type="Int32")
+
+@chdb_udf(return_type="Int32")
 def mul_udf(lhs, rhs):
     return int(lhs) * int(rhs)
+
+
+class TestUDF(unittest.TestCase):
+    def test_sum_udf(self):
+        ret = query("select sum_udf(12,22)", "Debug")
+        self.assertEqual(str(ret), '"34"\n')
+
+    def test_return_Int32(self):
+        ret = query("select mul_udf(12,22) + 1", "Debug")
+        self.assertEqual(str(ret), "265\n")
+
+    def test_define_in_function(self):
+        @chdb_udf()
+        def sum_udf2(lhs, rhs):
+            return int(lhs) + int(rhs)
+
+        ret = query("select sum_udf2(11, 22)", "Debug")
+        self.assertEqual(str(ret), '"33"\n')
+
+
+class TestUDFinSession(unittest.TestCase):
+    def test_sum_udf(self):
+        with Session() as session:
+            ret = session.query("select sum_udf(12,22)", "Debug")
+            self.assertEqual(str(ret), '"34"\n')
+
+    def test_return_Int32(self):
+        with Session() as session:
+            ret = session.query("select mul_udf(12,22) + 1", "Debug")
+            self.assertEqual(str(ret), "265\n")
+
+    def test_define_in_function(self):
+        @chdb_udf()
+        def sum_udf2(lhs, rhs):
+            return int(lhs) + int(rhs)
+
+        with Session() as session:
+            ret = session.query("select sum_udf2(11, 22)", "Debug")
+            self.assertEqual(str(ret), '"33"\n')
+
+if __name__ == "__main__":
+    unittest.main()