Fixed python query runner and add unittest for python query runner (#4731)

* add test for python runner

* fixed CustomPrint

* remove print

* Reformatted to pass our backend lint tests

Reformatted using:

  $ make format

---------

Co-authored-by: Justin Clift <justin@postgresql.org>

tests/query_runner/test_python.py

@@ -1,8 +1,77 @@
+from datetime import datetime
 from unittest import TestCase
 
+import mock
+
 from redash.query_runner.python import Python
 
 
+class TestPythonQueryRunner(TestCase):
+    def setUp(self):
+        self.python = Python({})
+
+    @mock.patch("datetime.datetime")
+    def test_print_in_query_string_success(self, mock_dt):
+        query_string = "print('test')"
+        mock_dt.utcnow = mock.Mock(return_value=datetime(1901, 12, 21))
+        result = self.python.run_query(query_string, "user")
+        self.assertEqual(result[0], '{"rows": [], "columns": [], "log": ["[1901-12-21T00:00:00] test"]}')
+
+    def test_empty_result(self):
+        query_string = "result={}"
+        result = self.python.run_query(query_string, "user")
+        self.assertEqual(result[0], '{"log": []}')
+
+    def test_invalidate_result_type_string(self):
+        query_string = "result='string'"
+        result = self.python.run_query(query_string, "user")
+        self.assertEqual(result[0], None)
+
+    def test_invalidate_result_type_int(self):
+        query_string = "result=100"
+        result = self.python.run_query(query_string, "user")
+        self.assertEqual(result[0], None)
+
+    def test_validate_result_type(self):
+        query_string = (
+            "result="
+            '{"columns": [{"name": "col1", "type": TYPE_STRING},'
+            '{"name": "col2", "type": TYPE_INTEGER}],'
+            '"rows": [{"col1": "foo", "col2": 100},'
+            '{"col1": "bar", "col2": 200}]}'
+        )
+        result = self.python.run_query(query_string, "user")
+        self.assertEqual(
+            result[0],
+            '{"columns": [{"name": "col1", "type": "string"},'
+            ' {"name": "col2", "type": "integer"}],'
+            ' "rows": [{"col1": "foo", "col2": 100},'
+            ' {"col1": "bar", "col2": 200}],'
+            ' "log": []}',
+        )
+
+    @mock.patch("datetime.datetime")
+    def test_validate_result_type_with_print(self, mock_dt):
+        mock_dt.utcnow = mock.Mock(return_value=datetime(1901, 12, 21))
+        query_string = (
+            'print("test")\n'
+            "result="
+            '{"columns": [{"name": "col1", "type": TYPE_STRING},'
+            '{"name": "col2", "type": TYPE_INTEGER}],'
+            '"rows": [{"col1": "foo", "col2": 100},'
+            '{"col1": "bar", "col2": 200}]}'
+        )
+        result = self.python.run_query(query_string, "user")
+        self.assertEqual(
+            result[0],
+            '{"columns": [{"name": "col1", "type": "string"},'
+            ' {"name": "col2", "type": "integer"}],'
+            ' "rows": [{"col1": "foo", "col2": 100},'
+            ' {"col1": "bar", "col2": 200}],'
+            ' "log": ["[1901-12-21T00:00:00] test"]}',
+        )
+
+
 class TestPython(TestCase):
     def test_sorted_safe_builtins(self):
         src = list(Python.safe_builtins)