[AIRFLOW-5755] Fixed most problems with py27

Author: potiuk

airflow/models/dag.py

@@ -1148,7 +1148,10 @@ def pickle(self, session=None):
     def tree_view(self):
         """Print an ASCII tree representation of the DAG."""
         def get_downstream(task, level=0):
-            print((" " * level * 4) + str(task))
+            line = (" " * level * 4) + str(task)
+            if six.PY2:
+                line = line.decode("utf-8")
+            print(line)
             level += 1
             for t in task.downstream_list:
                 get_downstream(t, level)

tests/contrib/hooks/test_ssh_hook.py

@@ -211,14 +211,15 @@ def test_ssh_connection_old_cm(self):
     def test_tunnel(self):
         hook = SSHHook(ssh_conn_id='ssh_default')
 
+        import psutil
         import subprocess
         import socket
 
         subprocess_kwargs = dict(
             args=["python", "-c", HELLO_SERVER_CMD],
             stdout=subprocess.PIPE,
         )
-        with subprocess.Popen(**subprocess_kwargs) as server_handle, hook.create_tunnel(2135, 2134):
+        with psutil.Popen(**subprocess_kwargs) as server_handle, hook.create_tunnel(2135, 2134):
             server_output = server_handle.stdout.read(5)
             self.assertEqual(b"ready", server_output)
             socket = socket.socket()

tests/models/test_baseoperator.py

@@ -19,7 +19,7 @@
 
 import datetime
 import unittest
-from unittest import mock
+from tests.compat import mock
 import uuid
 from collections import namedtuple
 
@@ -214,12 +214,11 @@ def test_nested_template_fields_declared_must_exist(self):
         with DAG("test-dag", start_date=DEFAULT_DATE):
             task = DummyOperator(task_id="op1")
 
-        with self.assertRaises(AttributeError) as e:
+        re = "('ClassWithCustomAttributes' object|ClassWithCustomAttributes instance) " \
+             "has no attribute 'missing_field'"
+        with self.assertRaisesRegexp(AttributeError, re):
             task.render_template(ClassWithCustomAttributes(template_fields=["missing_field"]), {})
 
-        self.assertEqual("'ClassWithCustomAttributes' object has no attribute 'missing_field'",
-                         str(e.exception))
-
     def test_jinja_invalid_expression_is_just_propagated(self):
         """Test render_template propagates Jinja invalid expression errors."""
         with DAG("test-dag", start_date=DEFAULT_DATE):

tests/models/test_dag.py

@@ -815,7 +815,7 @@ def test_roots(self):
             t5 = DummyOperator(task_id="t5")
             [t1, t2] >> t3 >> [t4, t5]
 
-            self.assertCountEqual(dag.roots, [t1, t2])
+            six.assertCountEqual(self, dag.roots, [t1, t2])
 
     def test_leaves(self):
         """Verify if dag.leaves returns the leaf tasks of a DAG."""
@@ -827,7 +827,7 @@ def test_leaves(self):
             t5 = DummyOperator(task_id="t5")
             [t1, t2] >> t3 >> [t4, t5]
 
-            self.assertCountEqual(dag.leaves, [t4, t5])
+            six.assertCountEqual(self, dag.leaves, [t4, t5])
 
     def test_tree_view(self):
         """Verify correctness of dag.tree_view()."""

tests/test_sentry.py

@@ -19,7 +19,7 @@
 
 import datetime
 import unittest
-from unittest.mock import Mock, MagicMock
+from tests.compat import Mock, MagicMock
 from freezegun import freeze_time
 
 from sentry_sdk import configure_scope

tests/test_utils/system_tests_class.py

@@ -17,7 +17,11 @@
 # specific language governing permissions and limitations
 # under the License.
 import os
-from contextlib import ContextDecorator
+try:
+    from contextdecorator import ContextDecorator
+except ImportError:
+    from contextlib import ContextDecorator
+
 from shutil import move
 from tempfile import mkdtemp
 from unittest import TestCase, skip

tests/www/test_views.py

@@ -557,14 +557,14 @@ def test_import_variables(self):
             self.assertEqual(case_b_dict, db_dict['dict_key'])
 
 
-@conf_vars({("webserver", "base_url"): "http://localhost:8080/test"})
 class TestMountPoint(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
         # Clear cached app to remount base_url forcefully
         application.app = None
-        app = application.cached_app(config={'WTF_CSRF_ENABLED': False}, testing=True)
-        cls.client = Client(app, BaseResponse)
+        with conf_vars({("webserver", "base_url"): "http://localhost:8080/test"}):
+            app = application.cached_app(config={'WTF_CSRF_ENABLED': False}, testing=True)
+            cls.client = Client(app, BaseResponse)
 
     @classmethod
     def tearDownClass(cls):