From 4b813f2460f7620f1ec457d1f1b7575f7c14fc94 Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Mon, 19 Sep 2022 14:27:21 +0200
Subject: [PATCH] Adjust pylint plugin for tests directory (#78727)

* Add module_name to parametrize

* Add tests for tests directory

* Apply patch from mib1185

* Adjust plugin to allow imports from component being tested
---
 pylint/plugins/hass_imports.py | 17 ++++++++++
 tests/pylint/test_imports.py   | 59 ++++++++++++++++++++++++++++------
 2 files changed, 67 insertions(+), 9 deletions(-)

diff --git a/pylint/plugins/hass_imports.py b/pylint/plugins/hass_imports.py
index 4b57794a07f..718f9550312 100644
--- a/pylint/plugins/hass_imports.py
+++ b/pylint/plugins/hass_imports.py
@@ -333,12 +333,22 @@ class HassImportsFormatChecker(BaseChecker):  # type: ignore[misc]
 
     def visit_import(self, node: nodes.Import) -> None:
         """Called when a Import node is visited."""
+        if self.current_package is None:
+            return
         for module, _alias in node.names:
             if module.startswith(f"{self.current_package}."):
                 self.add_message("hass-relative-import", node=node)
+                continue
             if module.startswith("homeassistant.components.") and module.endswith(
                 "const"
             ):
+                if (
+                    self.current_package.startswith("tests.components.")
+                    and self.current_package.split(".")[2] == module.split(".")[2]
+                ):
+                    # Ignore check if the component being tested matches
+                    # the component being imported from
+                    continue
                 self.add_message("hass-component-root-import", node=node)
 
     def _visit_importfrom_relative(
@@ -389,6 +399,13 @@ class HassImportsFormatChecker(BaseChecker):  # type: ignore[misc]
             node.modname.endswith(".const")
             or "const" in {names[0] for names in node.names}
         ):
+            if (
+                self.current_package.startswith("tests.components.")
+                and self.current_package.split(".")[2] == node.modname.split(".")[2]
+            ):
+                # Ignore check if the component being tested matches
+                # the component being imported from
+                return
             self.add_message("hass-component-root-import", node=node)
             return
         if obsolete_imports := _OBSOLETE_IMPORT.get(node.modname):
diff --git a/tests/pylint/test_imports.py b/tests/pylint/test_imports.py
index 130947ac68d..5c8bad28902 100644
--- a/tests/pylint/test_imports.py
+++ b/tests/pylint/test_imports.py
@@ -148,22 +148,41 @@ def test_bad_import(
 
 
 @pytest.mark.parametrize(
-    "import_node",
+    "import_node,module_name",
     [
-        "from homeassistant.components import climate",
-        "from homeassistant.components.climate import ClimateEntityFeature",
+        (
+            "from homeassistant.components import climate",
+            "homeassistant.components.pylint_test.climate",
+        ),
+        (
+            "from homeassistant.components.climate import ClimateEntityFeature",
+            "homeassistant.components.pylint_test.climate",
+        ),
+        (
+            "from homeassistant.components.pylint_test import const",
+            "tests.components.pylint_test.climate",
+        ),
+        (
+            "from homeassistant.components.pylint_test.const import CONSTANT",
+            "tests.components.pylint_test.climate",
+        ),
+        (
+            "import homeassistant.components.pylint_test.const as climate",
+            "tests.components.pylint_test.climate",
+        ),
     ],
 )
 def test_good_root_import(
     linter: UnittestLinter,
     imports_checker: BaseChecker,
     import_node: str,
+    module_name: str,
 ) -> None:
     """Ensure bad root imports are rejected."""
 
     node = astroid.extract_node(
         f"{import_node} #@",
-        "homeassistant.components.pylint_test.climate",
+        module_name,
     )
     imports_checker.visit_module(node.parent)
 
@@ -175,23 +194,45 @@ def test_good_root_import(
 
 
 @pytest.mark.parametrize(
-    "import_node",
+    "import_node,module_name",
     [
-        "import homeassistant.components.climate.const as climate",
-        "from homeassistant.components.climate import const",
-        "from homeassistant.components.climate.const import ClimateEntityFeature",
+        (
+            "import homeassistant.components.climate.const as climate",
+            "homeassistant.components.pylint_test.climate",
+        ),
+        (
+            "from homeassistant.components.climate import const",
+            "homeassistant.components.pylint_test.climate",
+        ),
+        (
+            "from homeassistant.components.climate.const import ClimateEntityFeature",
+            "homeassistant.components.pylint_test.climate",
+        ),
+        (
+            "from homeassistant.components.climate import const",
+            "tests.components.pylint_test.climate",
+        ),
+        (
+            "from homeassistant.components.climate.const import CONSTANT",
+            "tests.components.pylint_test.climate",
+        ),
+        (
+            "import homeassistant.components.climate.const as climate",
+            "tests.components.pylint_test.climate",
+        ),
     ],
 )
 def test_bad_root_import(
     linter: UnittestLinter,
     imports_checker: BaseChecker,
     import_node: str,
+    module_name: str,
 ) -> None:
     """Ensure bad root imports are rejected."""
 
     node = astroid.extract_node(
         f"{import_node} #@",
-        "homeassistant.components.pylint_test.climate",
+        module_name,
     )
     imports_checker.visit_module(node.parent)