[airflow] 05/06: Create AirflowMacroPluginRemoved rule and tests (#11285)

[ka...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

kaxilnaik pushed a commit to branch v1-10-test
in repository https://gitbox.apache.org/repos/asf/airflow.git

commit 550d9dda97e55de94265d0a1af8bd3e906a8254a
Author: Jacob Hoffman <72...@users.noreply.github.com>
AuthorDate: Sat Oct 10 03:57:02 2020 -0700

    Create AirflowMacroPluginRemoved rule and tests (#11285)
---
 .../upgrade/rules/airflow_macro_plugin_removed.py  | 53 +++++++++++++++++
 .../rules/test_airflow_macro_plugin_removed.py     | 68 ++++++++++++++++++++++
 2 files changed, 121 insertions(+)

diff --git a/airflow/upgrade/rules/airflow_macro_plugin_removed.py b/airflow/upgrade/rules/airflow_macro_plugin_removed.py
new file mode 100644
index 0000000..ca4f546
--- /dev/null
+++ b/airflow/upgrade/rules/airflow_macro_plugin_removed.py
@@ -0,0 +1,53 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import absolute_import
+
+from airflow import conf
+from airflow.upgrade.rules.base_rule import BaseRule
+from airflow.utils.dag_processing import list_py_file_paths
+
+
+class AirflowMacroPluginRemovedRule(BaseRule):
+
+    title = "Remove airflow.AirflowMacroPlugin class"
+
+    description = "The airflow.AirflowMacroPlugin class has been removed."
+
+    MACRO_PLUGIN_CLASS = "airflow.AirflowMacroPlugin"
+
+    def _change_info(self, file_path, line_number):
+        return "{} will be removed. Affected file: {} (line {})".format(
+            self.MACRO_PLUGIN_CLASS, file_path, line_number
+        )
+
+    def _check_file(self, file_path):
+        problems = []
+        class_name_to_check = self.MACRO_PLUGIN_CLASS.split(".")[-1]
+        with open(file_path, "r") as file_pointer:
+            for line_number, line in enumerate(file_pointer, 1):
+                if class_name_to_check in line:
+                    problems.append(self._change_info(file_path, line_number))
+        return problems
+
+    def check(self):
+        dag_folder = conf.get("core", "dags_folder")
+        file_paths = list_py_file_paths(directory=dag_folder, include_examples=False)
+        problems = []
+        for file_path in file_paths:
+            problems.extend(self._check_file(file_path))
+        return problems
diff --git a/tests/upgrade/rules/test_airflow_macro_plugin_removed.py b/tests/upgrade/rules/test_airflow_macro_plugin_removed.py
new file mode 100644
index 0000000..2d7b7ba
--- /dev/null
+++ b/tests/upgrade/rules/test_airflow_macro_plugin_removed.py
@@ -0,0 +1,68 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from contextlib import contextmanager
+from unittest import TestCase
+
+from tempfile import NamedTemporaryFile
+from tests.compat import mock
+
+from airflow.upgrade.rules.airflow_macro_plugin_removed import AirflowMacroPluginRemovedRule
+
+
+@contextmanager
+def create_temp_file(mock_list_files, lines):
+    with NamedTemporaryFile("w+") as temp_file:
+        mock_list_files.return_value = [temp_file.name]
+        temp_file.writelines("\n".join(lines))
+        temp_file.flush()
+        yield temp_file
+
+
+@mock.patch("airflow.upgrade.rules.airflow_macro_plugin_removed.list_py_file_paths")
+class TestAirflowMacroPluginRemovedRule(TestCase):
+    def test_valid_check(self, mock_list_files):
+        lines = ["import foo.bar.baz"]
+        with create_temp_file(mock_list_files, lines):
+            rule = AirflowMacroPluginRemovedRule()
+            assert isinstance(rule.description, str)
+            assert isinstance(rule.title, str)
+
+            msgs = rule.check()
+            assert 0 == len(msgs)
+
+    def test_invalid_check(self, mock_list_files):
+        lines = [
+            "import airflow.AirflowMacroPlugin",
+            "from airflow import AirflowMacroPlugin",
+        ]
+        with create_temp_file(mock_list_files, lines) as temp_file:
+
+            rule = AirflowMacroPluginRemovedRule()
+
+            assert isinstance(rule.description, str)
+            assert isinstance(rule.title, str)
+
+            msgs = rule.check()
+            assert 2 == len(msgs)
+
+            base_message = "airflow.AirflowMacroPlugin will be removed. Affected file: {}".format(
+                temp_file.name
+            )
+            expected_messages = [
+                "{} (line {})".format(base_message, line_number) for line_number in [1, 2]
+            ]
+            assert expected_messages == msgs