Adjust relative-import plugin for tests (home-assistant#78742)

johntdyer · Sep 19, 2022 · 0dcbc85 · 0dcbc85
1 parent cd6959d
commit 0dcbc85
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 14 deletions.
diff --git a/pylint/plugins/hass_imports.py b/pylint/plugins/hass_imports.py
@@ -345,8 +345,10 @@ def _visit_importfrom_relative(
         self, current_package: str, node: nodes.ImportFrom
     ) -> None:
         """Called when a ImportFrom node is visited."""
-        if node.level <= 1 or not current_package.startswith(
-            "homeassistant.components"
+        if (
+            node.level <= 1
+            or not current_package.startswith("homeassistant.components.")
+            and not current_package.startswith("tests.components.")
         ):
             return
         split_package = current_package.split(".")
@@ -372,18 +374,17 @@ def visit_importfrom(self, node: nodes.ImportFrom) -> None:
         ):
             self.add_message("hass-relative-import", node=node)
             return
-        if self.current_package.startswith("homeassistant.components."):
-            current_component = self.current_package.split(".")[2]
-            if node.modname == "homeassistant.components":
-                for name in node.names:
-                    if name[0] == current_component:
-                        self.add_message("hass-relative-import", node=node)
-                return
-            if node.modname.startswith(
-                f"homeassistant.components.{current_component}."
-            ):
-                self.add_message("hass-relative-import", node=node)
-                return
+        for root in ("homeassistant", "tests"):
+            if self.current_package.startswith(f"{root}.components."):
+                current_component = self.current_package.split(".")[2]
+                if node.modname == f"{root}.components":
+                    for name in node.names:
+                        if name[0] == current_component:
+                            self.add_message("hass-relative-import", node=node)
+                    return
+                if node.modname.startswith(f"{root}.components.{current_component}."):
+                    self.add_message("hass-relative-import", node=node)
+                    return
         if node.modname.startswith("homeassistant.components.") and (
             node.modname.endswith(".const")
             or "const" in {names[0] for names in node.names}

diff --git a/tests/pylint/test_imports.py b/tests/pylint/test_imports.py
@@ -35,6 +35,7 @@
         ("homeassistant.components.pylint_test.api.hub", "..const", "CONSTANT"),
         ("homeassistant.components.pylint_test.api.hub", "..", "CONSTANT"),
         ("homeassistant.components.pylint_test.api.hub", "...", "pylint_test"),
+        ("tests.components.pylint_test.api.hub", "..const", "CONSTANT"),
     ],
 )
 def test_good_import(
@@ -101,6 +102,18 @@ def test_good_import(
             "CONSTANT",
             "hass-relative-import",
         ),
+        (
+            "tests.components.pylint_test.api.hub",
+            "tests.components.pylint_test.const",
+            "CONSTANT",
+            "hass-relative-import",
+        ),
+        (
+            "tests.components.pylint_test.api.hub",
+            "...const",
+            "CONSTANT",
+            "hass-absolute-import",
+        ),
     ],
 )
 def test_bad_import(