Detect references to deprecated classes in test_core_to_contrib.py (a…

…pache#9553)
cfei18 · Jun 29, 2020 · 6b18ed4 · 6b18ed4
1 parent 40add26
commit 6b18ed4
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 15 deletions.
diff --git a/airflow/operators/presto_check_operator.py b/airflow/operators/presto_check_operator.py
@@ -15,62 +15,62 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-"""This module is deprecated. Please use `airflow.operators.check_operator`."""
+"""This module is deprecated. Please use `airflow.operators.sql`."""
 
 import warnings
 
 # pylint: disable=unused-import
-from airflow.operators.check_operator import CheckOperator, IntervalCheckOperator, ValueCheckOperator  # noqa
+from airflow.operators.sql import SQLCheckOperator, SQLIntervalCheckOperator, SQLValueCheckOperator  # noqa
 
 warnings.warn(
-    "This module is deprecated. Please use `airflow.operators.check_operator`.",
+    "This module is deprecated. Please use `airflow.operators.sql`.",
     DeprecationWarning, stacklevel=2
 )
 
 
-class PrestoCheckOperator(CheckOperator):
+class PrestoCheckOperator(SQLCheckOperator):
     """
     This class is deprecated.
-    Please use `airflow.operators.check_operator.CheckOperator`.
+    Please use `airflow.operators.sql.SQLCheckOperator`.
     """
 
     def __init__(self, *args, **kwargs):
         warnings.warn(
             """This class is deprecated.
-            Please use `airflow.operators.check_operator.CheckOperator`.""",
+            Please use `airflow.operators.sql.SQLCheckOperator`.""",
             DeprecationWarning, stacklevel=3
         )
         super().__init__(*args, **kwargs)
 
 
-class PrestoIntervalCheckOperator(IntervalCheckOperator):
+class PrestoIntervalCheckOperator(SQLIntervalCheckOperator):
     """
     This class is deprecated.
-    Please use `airflow.operators.check_operator.IntervalCheckOperator`.
+    Please use `airflow.operators.sql.SQLIntervalCheckOperator`.
     """
 
     def __init__(self, *args, **kwargs):
         warnings.warn(
             """
             This class is deprecated.l
-            Please use `airflow.operators.check_operator.IntervalCheckOperator`.
+            Please use `airflow.operators.sql.SQLIntervalCheckOperator`.
             """,
             DeprecationWarning, stacklevel=3
         )
         super().__init__(*args, **kwargs)
 
 
-class PrestoValueCheckOperator(ValueCheckOperator):
+class PrestoValueCheckOperator(SQLValueCheckOperator):
     """
     This class is deprecated.
-    Please use `airflow.operators.check_operator.ValueCheckOperator`.
+    Please use `airflow.operators.sql.SQLValueCheckOperator`.
     """
 
     def __init__(self, *args, **kwargs):
         warnings.warn(
             """
             This class is deprecated.l
-            Please use `airflow.operators.check_operator.ValueCheckOperator`.
+            Please use `airflow.operators.sql.SQLValueCheckOperator`.
             """,
             DeprecationWarning, stacklevel=3
         )

diff --git a/tests/deprecated_classes.py b/tests/deprecated_classes.py
@@ -1148,15 +1148,15 @@
         'airflow.operators.papermill_operator.PapermillOperator',
     ),
     (
-        'airflow.operators.check_operator.CheckOperator',
+        'airflow.operators.sql.SQLCheckOperator',
         'airflow.operators.presto_check_operator.PrestoCheckOperator',
     ),
     (
-        'airflow.operators.check_operator.IntervalCheckOperator',
+        'airflow.operators.sql.SQLIntervalCheckOperator',
         'airflow.operators.presto_check_operator.PrestoIntervalCheckOperator',
     ),
     (
-        'airflow.operators.check_operator.ValueCheckOperator',
+        'airflow.operators.sql.SQLValueCheckOperator',
         'airflow.operators.presto_check_operator.PrestoValueCheckOperator',
     ),
     (

diff --git a/tests/test_core_to_contrib.py b/tests/test_core_to_contrib.py
@@ -102,3 +102,24 @@ def test_is_subclass(self, parent_class_path, sub_class_path):
     def test_warning_on_import(self, new_path, old_path):
         self.skip_test_with_mssql_in_py38(new_path, old_path)
         self.assert_proper_import(old_path, new_path)
+
+    def test_no_redirect_to_deprecated_classes(self):
+        """
+        When we have the following items:
+        new_A, old_B
+        old_B, old_C
+
+        This will tell us to use new_A instead of old_B.
+        """
+        all_classes_by_old = {
+            old: new for new, old in ALL
+        }
+
+        for new, old in ALL:
+            # Using if statement allows us to create a developer-friendly message only when we need it.
+            # Otherwise, it wouldn't always be possible - KeyError
+            if new in all_classes_by_old:
+                raise AssertionError(
+                    f'Deprecation "{old}" to "{new}" is incorrect. '
+                    f'Please use \"{all_classes_by_old[new]}\" instead of "{old}".'
+                )