From 950ee3451495b097d789ae60dccf6b61ff45ed51 Mon Sep 17 00:00:00 2001
From: Jan Bouwhuis <jbouwh@users.noreply.github.com>
Date: Mon, 13 Feb 2023 12:12:34 +0100
Subject: [PATCH] Correct typing on hass_storage fixure (#87986)

Co-authored-by: epenet <6771947+epenet@users.noreply.github.com>
---
 pylint/plugins/hass_enforce_type_hints.py |  1 +
 tests/common.py                           | 23 ++++++++++++++++++-----
 tests/components/mqtt/test_init.py        |  2 +-
 tests/conftest.py                         |  2 +-
 4 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/pylint/plugins/hass_enforce_type_hints.py b/pylint/plugins/hass_enforce_type_hints.py
index e1637d51cfe..99d48f165be 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] = {
     "hass_read_only_access_token": "str",
     "hass_read_only_user": "MockUser",
     "hass_recorder": "Callable[..., HomeAssistant]",
+    "hass_storage": "dict[str, Any]",
     "hass_supervisor_access_token": "str",
     "hass_supervisor_user": "MockUser",
     "hass_ws_client": "WebSocketGenerator",
diff --git a/tests/common.py b/tests/common.py
index eb8b5fe9753..e45efb028e6 100644
--- a/tests/common.py
+++ b/tests/common.py
@@ -3,7 +3,14 @@ from __future__ import annotations
 
 import asyncio
 from collections import OrderedDict
-from collections.abc import Awaitable, Callable, Collection, Mapping, Sequence
+from collections.abc import (
+    Awaitable,
+    Callable,
+    Collection,
+    Generator,
+    Mapping,
+    Sequence,
+)
 from contextlib import contextmanager
 from datetime import datetime, timedelta, timezone
 import functools as ft
@@ -1205,7 +1212,9 @@ class MockEntity(entity.Entity):
 
 
 @contextmanager
-def mock_storage(data=None):
+def mock_storage(
+    data: dict[str, Any] | None = None
+) -> Generator[dict[str, Any], None, None]:
     """Mock storage.
 
     Data is a dict {'key': {'version': version, 'data': data}}
@@ -1217,7 +1226,9 @@ def mock_storage(data=None):
 
     orig_load = storage.Store._async_load
 
-    async def mock_async_load(store):
+    async def mock_async_load(
+        store: storage.Store,
+    ) -> dict[str, Any] | list[Any] | None:
         """Mock version of load."""
         if store._data is None:
             # No data to load
@@ -1237,14 +1248,16 @@ def mock_storage(data=None):
         _LOGGER.debug("Loading data for %s: %s", store.key, loaded)
         return loaded
 
-    async def mock_write_data(store, path, data_to_write):
+    async def mock_write_data(
+        store: storage.Store, path: str, data_to_write: dict[str, Any]
+    ) -> None:
         """Mock version of write data."""
         # To ensure that the data can be serialized
         _LOGGER.debug("Writing data to %s: %s", store.key, data_to_write)
         raise_contains_mocks(data_to_write)
         data[store.key] = json.loads(json.dumps(data_to_write, cls=store._encoder))
 
-    async def mock_remove(store):
+    async def mock_remove(store: storage.Store) -> None:
         """Remove data."""
         data.pop(store.key, None)
 
diff --git a/tests/components/mqtt/test_init.py b/tests/components/mqtt/test_init.py
index 934a15acb46..37d7e409a8f 100644
--- a/tests/components/mqtt/test_init.py
+++ b/tests/components/mqtt/test_init.py
@@ -98,7 +98,7 @@ def sensor_platforms_only() -> Generator[None, None, None]:
 
 
 @pytest.fixture(autouse=True)
-def mock_storage(hass_storage: Generator[dict[str, Any], None, None]) -> None:
+def mock_storage(hass_storage: dict[str, Any]) -> None:
     """Autouse hass_storage for the TestCase tests."""
 
 
diff --git a/tests/conftest.py b/tests/conftest.py
index 510be346540..186b67dff6e 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -302,7 +302,7 @@ def bcrypt_cost() -> Generator[None, None, None]:
 
 
 @pytest.fixture
-def hass_storage():
+def hass_storage() -> Generator[dict[str, Any], None, None]:
     """Fixture to mock storage."""
     with mock_storage() as stored_data:
         yield stored_data