From 0974ea9a5a2af44d661728e47c71852b000b01f3 Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Fri, 31 May 2024 13:06:49 +0200
Subject: [PATCH] Adjust "hass" type hint for test fixtures in pylint plugin
 (#118548)

Adjust "hass" type hint in pylint plugin
---
 pylint/plugins/hass_enforce_type_hints.py | 21 ++++++-------
 tests/pylint/test_enforce_type_hints.py   | 36 +++++++++++------------
 2 files changed, 29 insertions(+), 28 deletions(-)

diff --git a/pylint/plugins/hass_enforce_type_hints.py b/pylint/plugins/hass_enforce_type_hints.py
index 16449e2e5a0..c6c6986060f 100644
--- a/pylint/plugins/hass_enforce_type_hints.py
+++ b/pylint/plugins/hass_enforce_type_hints.py
@@ -113,6 +113,7 @@ _TEST_FIXTURES: dict[str, list[str] | str] = {
     "entity_registry_enabled_by_default": "None",
     "event_loop": "AbstractEventLoop",
     "freezer": "FrozenDateTimeFactory",
+    "hass": "HomeAssistant",
     "hass_access_token": "str",
     "hass_admin_credential": "Credentials",
     "hass_admin_user": "MockUser",
@@ -3218,16 +3219,6 @@ class HassTypeHintChecker(BaseChecker):
         if self._ignore_function(node, annotations):
             return
 
-        # Check that common arguments are correctly typed.
-        for arg_name, expected_type in _COMMON_ARGUMENTS.items():
-            arg_node, annotation = _get_named_annotation(node, arg_name)
-            if arg_node and not _is_valid_type(expected_type, annotation):
-                self.add_message(
-                    "hass-argument-type",
-                    node=arg_node,
-                    args=(arg_name, expected_type, node.name),
-                )
-
         # Check method or function matchers.
         if node.is_method():
             matchers = _METHOD_MATCH
@@ -3246,6 +3237,16 @@ class HassTypeHintChecker(BaseChecker):
                     return
             matchers = self._function_matchers
 
+        # Check that common arguments are correctly typed.
+        for arg_name, expected_type in _COMMON_ARGUMENTS.items():
+            arg_node, annotation = _get_named_annotation(node, arg_name)
+            if arg_node and not _is_valid_type(expected_type, annotation):
+                self.add_message(
+                    "hass-argument-type",
+                    node=arg_node,
+                    args=(arg_name, expected_type, node.name),
+                )
+
         for match in matchers:
             if not match.need_to_check_function(node):
                 continue
diff --git a/tests/pylint/test_enforce_type_hints.py b/tests/pylint/test_enforce_type_hints.py
index 68e1e14a34f..9f0f4905dab 100644
--- a/tests/pylint/test_enforce_type_hints.py
+++ b/tests/pylint/test_enforce_type_hints.py
@@ -1174,15 +1174,6 @@ def test_pytest_invalid_function(
 
     with assert_adds_messages(
         linter,
-        pylint.testutils.MessageTest(
-            msg_id="hass-argument-type",
-            node=hass_node,
-            args=("hass", ["HomeAssistant", "HomeAssistant | None"], "test_sample"),
-            line=3,
-            col_offset=4,
-            end_line=3,
-            end_col_offset=19,
-        ),
         pylint.testutils.MessageTest(
             msg_id="hass-return-type",
             node=func_node,
@@ -1228,6 +1219,15 @@ def test_pytest_invalid_function(
             end_line=6,
             end_col_offset=36,
         ),
+        pylint.testutils.MessageTest(
+            msg_id="hass-argument-type",
+            node=hass_node,
+            args=("hass", "HomeAssistant", "test_sample"),
+            line=3,
+            col_offset=4,
+            end_line=3,
+            end_col_offset=19,
+        ),
     ):
         type_hint_checker.visit_asyncfunctiondef(func_node)
 
@@ -1281,15 +1281,6 @@ def test_pytest_invalid_fixture(
 
     with assert_adds_messages(
         linter,
-        pylint.testutils.MessageTest(
-            msg_id="hass-argument-type",
-            node=hass_node,
-            args=("hass", ["HomeAssistant", "HomeAssistant | None"], "sample_fixture"),
-            line=6,
-            col_offset=4,
-            end_line=6,
-            end_col_offset=19,
-        ),
         pylint.testutils.MessageTest(
             msg_id="hass-argument-type",
             node=caplog_node,
@@ -1308,6 +1299,15 @@ def test_pytest_invalid_fixture(
             end_line=8,
             end_col_offset=29,
         ),
+        pylint.testutils.MessageTest(
+            msg_id="hass-argument-type",
+            node=hass_node,
+            args=("hass", "HomeAssistant", "sample_fixture"),
+            line=6,
+            col_offset=4,
+            end_line=6,
+            end_col_offset=19,
+        ),
     ):
         type_hint_checker.visit_asyncfunctiondef(func_node)