From 19e272a901ea670adf05c56cf2db68f1707f5d90 Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Mon, 8 Jul 2024 13:31:05 +0200
Subject: [PATCH] Improve type hints in iotawatt tests (#121505)

---
 tests/components/iotawatt/conftest.py    |  8 +++++---
 tests/components/iotawatt/test_init.py   | 14 +++++++++++---
 tests/components/iotawatt/test_sensor.py |  5 +++--
 3 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/tests/components/iotawatt/conftest.py b/tests/components/iotawatt/conftest.py
index f3a60e69021..9380154b53e 100644
--- a/tests/components/iotawatt/conftest.py
+++ b/tests/components/iotawatt/conftest.py
@@ -1,16 +1,18 @@
 """Test fixtures for IoTaWatt."""
 
-from unittest.mock import AsyncMock, patch
+from collections.abc import Generator
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 
 from homeassistant.components.iotawatt import DOMAIN
+from homeassistant.core import HomeAssistant
 
 from tests.common import MockConfigEntry
 
 
 @pytest.fixture
-def entry(hass):
+def entry(hass: HomeAssistant) -> MockConfigEntry:
     """Mock config entry added to HA."""
     entry = MockConfigEntry(domain=DOMAIN, data={"host": "1.2.3.4"})
     entry.add_to_hass(hass)
@@ -18,7 +20,7 @@ def entry(hass):
 
 
 @pytest.fixture
-def mock_iotawatt(entry):
+def mock_iotawatt(entry: MockConfigEntry) -> Generator[MagicMock]:
     """Mock iotawatt."""
     with patch("homeassistant.components.iotawatt.coordinator.Iotawatt") as mock:
         instance = mock.return_value
diff --git a/tests/components/iotawatt/test_init.py b/tests/components/iotawatt/test_init.py
index 8b707780eb4..de3a2f9f829 100644
--- a/tests/components/iotawatt/test_init.py
+++ b/tests/components/iotawatt/test_init.py
@@ -1,5 +1,7 @@
 """Test init."""
 
+from unittest.mock import MagicMock
+
 import httpx
 
 from homeassistant.config_entries import ConfigEntryState
@@ -8,8 +10,12 @@ from homeassistant.setup import async_setup_component
 
 from . import INPUT_SENSOR
 
+from tests.common import MockConfigEntry
 
-async def test_setup_unload(hass: HomeAssistant, mock_iotawatt, entry) -> None:
+
+async def test_setup_unload(
+    hass: HomeAssistant, mock_iotawatt: MagicMock, entry: MockConfigEntry
+) -> None:
     """Test we can setup and unload an entry."""
     mock_iotawatt.getSensors.return_value["sensors"]["my_sensor_key"] = INPUT_SENSOR
     assert await async_setup_component(hass, "iotawatt", {})
@@ -18,7 +24,7 @@ async def test_setup_unload(hass: HomeAssistant, mock_iotawatt, entry) -> None:
 
 
 async def test_setup_connection_failed(
-    hass: HomeAssistant, mock_iotawatt, entry
+    hass: HomeAssistant, mock_iotawatt: MagicMock, entry: MockConfigEntry
 ) -> None:
     """Test connection error during startup."""
     mock_iotawatt.connect.side_effect = httpx.ConnectError("")
@@ -27,7 +33,9 @@ async def test_setup_connection_failed(
     assert entry.state is ConfigEntryState.SETUP_RETRY
 
 
-async def test_setup_auth_failed(hass: HomeAssistant, mock_iotawatt, entry) -> None:
+async def test_setup_auth_failed(
+    hass: HomeAssistant, mock_iotawatt: MagicMock, entry: MockConfigEntry
+) -> None:
     """Test auth error during startup."""
     mock_iotawatt.connect.return_value = False
     assert await async_setup_component(hass, "iotawatt", {})
diff --git a/tests/components/iotawatt/test_sensor.py b/tests/components/iotawatt/test_sensor.py
index ecf2f97c67a..eb1a240a82f 100644
--- a/tests/components/iotawatt/test_sensor.py
+++ b/tests/components/iotawatt/test_sensor.py
@@ -1,6 +1,7 @@
 """Test setting up sensors."""
 
 from datetime import timedelta
+from unittest.mock import MagicMock
 
 from freezegun.api import FrozenDateTimeFactory
 
@@ -25,7 +26,7 @@ from tests.common import async_fire_time_changed
 
 
 async def test_sensor_type_input(
-    hass: HomeAssistant, freezer: FrozenDateTimeFactory, mock_iotawatt
+    hass: HomeAssistant, freezer: FrozenDateTimeFactory, mock_iotawatt: MagicMock
 ) -> None:
     """Test input sensors work."""
     assert await async_setup_component(hass, "iotawatt", {})
@@ -60,7 +61,7 @@ async def test_sensor_type_input(
 
 
 async def test_sensor_type_output(
-    hass: HomeAssistant, freezer: FrozenDateTimeFactory, mock_iotawatt
+    hass: HomeAssistant, freezer: FrozenDateTimeFactory, mock_iotawatt: MagicMock
 ) -> None:
     """Tests the sensor type of Output."""
     mock_iotawatt.getSensors.return_value["sensors"]["my_watthour_sensor_key"] = (