From a049d0e8464e08f4087d8680bb8f30bf7d3be20b Mon Sep 17 00:00:00 2001
From: Adam Goode <agoode@google.com>
Date: Sun, 3 Mar 2024 23:13:35 -0500
Subject: [PATCH] Add types throughout the prometheus tests (#112156)

---
 tests/components/prometheus/test_init.py | 132 ++++++++++++++++-------
 1 file changed, 93 insertions(+), 39 deletions(-)

diff --git a/tests/components/prometheus/test_init.py b/tests/components/prometheus/test_init.py
index 7ee534f91ce..812a88db16e 100644
--- a/tests/components/prometheus/test_init.py
+++ b/tests/components/prometheus/test_init.py
@@ -65,6 +65,8 @@ from homeassistant.helpers import entity_registry as er
 from homeassistant.setup import async_setup_component
 from homeassistant.util import dt as dt_util
 
+from tests.typing import ClientSessionGenerator
+
 PROMETHEUS_PATH = "homeassistant.components.prometheus"
 
 
@@ -77,7 +79,11 @@ class FilterTest:
 
 
 @pytest.fixture(name="client")
-async def setup_prometheus_client(hass, hass_client, namespace):
+async def setup_prometheus_client(
+    hass: HomeAssistant,
+    hass_client: ClientSessionGenerator,
+    namespace: str,
+):
     """Initialize an hass_client with Prometheus component."""
     # Reset registry
     prometheus_client.REGISTRY = prometheus_client.CollectorRegistry(auto_describe=True)
@@ -110,7 +116,12 @@ async def generate_latest_metrics(client):
 
 
 @pytest.mark.parametrize("namespace", [""])
-async def test_setup_enumeration(hass, hass_client, entity_registry, namespace):
+async def test_setup_enumeration(
+    hass: HomeAssistant,
+    hass_client: ClientSessionGenerator,
+    entity_registry: er.EntityRegistry,
+    namespace: str,
+) -> None:
     """Test that setup enumerates existing states/entities."""
 
     # The order of when things are created must be carefully controlled in
@@ -138,7 +149,9 @@ async def test_setup_enumeration(hass, hass_client, entity_registry, namespace):
 
 
 @pytest.mark.parametrize("namespace", [""])
-async def test_view_empty_namespace(client, sensor_entities) -> None:
+async def test_view_empty_namespace(
+    client: ClientSessionGenerator, sensor_entities: dict[str, er.RegistryEntry]
+) -> None:
     """Test prometheus metrics view."""
     body = await generate_latest_metrics(client)
 
@@ -162,7 +175,9 @@ async def test_view_empty_namespace(client, sensor_entities) -> None:
 
 
 @pytest.mark.parametrize("namespace", [None])
-async def test_view_default_namespace(client, sensor_entities) -> None:
+async def test_view_default_namespace(
+    client: ClientSessionGenerator, sensor_entities: dict[str, er.RegistryEntry]
+) -> None:
     """Test prometheus metrics view."""
     body = await generate_latest_metrics(client)
 
@@ -180,7 +195,9 @@ async def test_view_default_namespace(client, sensor_entities) -> None:
 
 
 @pytest.mark.parametrize("namespace", [""])
-async def test_sensor_unit(client, sensor_entities) -> None:
+async def test_sensor_unit(
+    client: ClientSessionGenerator, sensor_entities: dict[str, er.RegistryEntry]
+) -> None:
     """Test prometheus metrics for sensors with a unit."""
     body = await generate_latest_metrics(client)
 
@@ -210,7 +227,9 @@ async def test_sensor_unit(client, sensor_entities) -> None:
 
 
 @pytest.mark.parametrize("namespace", [""])
-async def test_sensor_without_unit(client, sensor_entities) -> None:
+async def test_sensor_without_unit(
+    client: ClientSessionGenerator, sensor_entities: dict[str, er.RegistryEntry]
+) -> None:
     """Test prometheus metrics for sensors without a unit."""
     body = await generate_latest_metrics(client)
 
@@ -234,7 +253,9 @@ async def test_sensor_without_unit(client, sensor_entities) -> None:
 
 
 @pytest.mark.parametrize("namespace", [""])
-async def test_sensor_device_class(client, sensor_entities) -> None:
+async def test_sensor_device_class(
+    client: ClientSessionGenerator, sensor_entities: dict[str, er.RegistryEntry]
+) -> None:
     """Test prometheus metrics for sensor with a device_class."""
     body = await generate_latest_metrics(client)
 
@@ -270,7 +291,9 @@ async def test_sensor_device_class(client, sensor_entities) -> None:
 
 
 @pytest.mark.parametrize("namespace", [""])
-async def test_input_number(client, input_number_entities) -> None:
+async def test_input_number(
+    client: ClientSessionGenerator, input_number_entities: dict[str, er.RegistryEntry]
+) -> None:
     """Test prometheus metrics for input_number."""
     body = await generate_latest_metrics(client)
 
@@ -294,7 +317,9 @@ async def test_input_number(client, input_number_entities) -> None:
 
 
 @pytest.mark.parametrize("namespace", [""])
-async def test_number(client, number_entities) -> None:
+async def test_number(
+    client: ClientSessionGenerator, number_entities: dict[str, er.RegistryEntry]
+) -> None:
     """Test prometheus metrics for number."""
     body = await generate_latest_metrics(client)
 
@@ -318,7 +343,9 @@ async def test_number(client, number_entities) -> None:
 
 
 @pytest.mark.parametrize("namespace", [""])
-async def test_battery(client, sensor_entities) -> None:
+async def test_battery(
+    client: ClientSessionGenerator, sensor_entities: dict[str, er.RegistryEntry]
+) -> None:
     """Test prometheus metrics for battery."""
     body = await generate_latest_metrics(client)
 
@@ -330,7 +357,10 @@ async def test_battery(client, sensor_entities) -> None:
 
 
 @pytest.mark.parametrize("namespace", [""])
-async def test_climate(client, climate_entities) -> None:
+async def test_climate(
+    client: ClientSessionGenerator,
+    climate_entities: dict[str, er.RegistryEntry | dict[str, Any]],
+) -> None:
     """Test prometheus metrics for climate entities."""
     body = await generate_latest_metrics(client)
 
@@ -366,7 +396,10 @@ async def test_climate(client, climate_entities) -> None:
 
 
 @pytest.mark.parametrize("namespace", [""])
-async def test_humidifier(client, humidifier_entities) -> None:
+async def test_humidifier(
+    client: ClientSessionGenerator,
+    humidifier_entities: dict[str, er.RegistryEntry | dict[str, Any]],
+) -> None:
     """Test prometheus metrics for humidifier entities."""
     body = await generate_latest_metrics(client)
 
@@ -397,7 +430,10 @@ async def test_humidifier(client, humidifier_entities) -> None:
 
 
 @pytest.mark.parametrize("namespace", [""])
-async def test_attributes(client, switch_entities) -> None:
+async def test_attributes(
+    client: ClientSessionGenerator,
+    switch_entities: dict[str, er.RegistryEntry | dict[str, Any]],
+) -> None:
     """Test prometheus metrics for entity attributes."""
     body = await generate_latest_metrics(client)
 
@@ -427,7 +463,9 @@ async def test_attributes(client, switch_entities) -> None:
 
 
 @pytest.mark.parametrize("namespace", [""])
-async def test_binary_sensor(client, binary_sensor_entities) -> None:
+async def test_binary_sensor(
+    client: ClientSessionGenerator, binary_sensor_entities: dict[str, er.RegistryEntry]
+) -> None:
     """Test prometheus metrics for binary_sensor."""
     body = await generate_latest_metrics(client)
 
@@ -445,7 +483,9 @@ async def test_binary_sensor(client, binary_sensor_entities) -> None:
 
 
 @pytest.mark.parametrize("namespace", [""])
-async def test_input_boolean(client, input_boolean_entities) -> None:
+async def test_input_boolean(
+    client: ClientSessionGenerator, input_boolean_entities: dict[str, er.RegistryEntry]
+) -> None:
     """Test prometheus metrics for input_boolean."""
     body = await generate_latest_metrics(client)
 
@@ -463,7 +503,9 @@ async def test_input_boolean(client, input_boolean_entities) -> None:
 
 
 @pytest.mark.parametrize("namespace", [""])
-async def test_light(client, light_entities) -> None:
+async def test_light(
+    client: ClientSessionGenerator, light_entities: dict[str, er.RegistryEntry]
+) -> None:
     """Test prometheus metrics for lights."""
     body = await generate_latest_metrics(client)
 
@@ -499,7 +541,9 @@ async def test_light(client, light_entities) -> None:
 
 
 @pytest.mark.parametrize("namespace", [""])
-async def test_lock(client, lock_entities) -> None:
+async def test_lock(
+    client: ClientSessionGenerator, lock_entities: dict[str, er.RegistryEntry]
+) -> None:
     """Test prometheus metrics for lock."""
     body = await generate_latest_metrics(client)
 
@@ -517,7 +561,9 @@ async def test_lock(client, lock_entities) -> None:
 
 
 @pytest.mark.parametrize("namespace", [""])
-async def test_cover(client, cover_entities) -> None:
+async def test_cover(
+    client: ClientSessionGenerator, cover_entities: dict[str, er.RegistryEntry]
+) -> None:
     """Test prometheus metrics for cover."""
     data = {**cover_entities}
     body = await generate_latest_metrics(client)
@@ -576,7 +622,9 @@ async def test_cover(client, cover_entities) -> None:
 
 
 @pytest.mark.parametrize("namespace", [""])
-async def test_device_tracker(client, device_tracker_entities) -> None:
+async def test_device_tracker(
+    client: ClientSessionGenerator, device_tracker_entities: dict[str, er.RegistryEntry]
+) -> None:
     """Test prometheus metrics for device_tracker."""
     body = await generate_latest_metrics(client)
 
@@ -593,7 +641,9 @@ async def test_device_tracker(client, device_tracker_entities) -> None:
 
 
 @pytest.mark.parametrize("namespace", [""])
-async def test_counter(client, counter_entities) -> None:
+async def test_counter(
+    client: ClientSessionGenerator, counter_entities: dict[str, er.RegistryEntry]
+) -> None:
     """Test prometheus metrics for counter."""
     body = await generate_latest_metrics(client)
 
@@ -605,7 +655,9 @@ async def test_counter(client, counter_entities) -> None:
 
 
 @pytest.mark.parametrize("namespace", [""])
-async def test_update(client, update_entities) -> None:
+async def test_update(
+    client: ClientSessionGenerator, update_entities: dict[str, er.RegistryEntry]
+) -> None:
     """Test prometheus metrics for update."""
     body = await generate_latest_metrics(client)
 
@@ -625,9 +677,9 @@ async def test_update(client, update_entities) -> None:
 async def test_renaming_entity_name(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,
-    client,
-    sensor_entities,
-    climate_entities,
+    client: ClientSessionGenerator,
+    sensor_entities: dict[str, er.RegistryEntry],
+    climate_entities: dict[str, er.RegistryEntry | dict[str, Any]],
 ) -> None:
     """Test renaming entity name."""
     data = {**sensor_entities, **climate_entities}
@@ -751,9 +803,9 @@ async def test_renaming_entity_name(
 async def test_renaming_entity_id(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,
-    client,
-    sensor_entities,
-    climate_entities,
+    client: ClientSessionGenerator,
+    sensor_entities: dict[str, er.RegistryEntry],
+    climate_entities: dict[str, er.RegistryEntry | dict[str, Any]],
 ) -> None:
     """Test renaming entity id."""
     data = {**sensor_entities, **climate_entities}
@@ -831,9 +883,9 @@ async def test_renaming_entity_id(
 async def test_deleting_entity(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,
-    client,
-    sensor_entities,
-    climate_entities,
+    client: ClientSessionGenerator,
+    sensor_entities: dict[str, er.RegistryEntry],
+    climate_entities: dict[str, er.RegistryEntry | dict[str, Any]],
 ) -> None:
     """Test deleting a entity."""
     data = {**sensor_entities, **climate_entities}
@@ -910,9 +962,9 @@ async def test_deleting_entity(
 async def test_disabling_entity(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,
-    client,
-    sensor_entities,
-    climate_entities,
+    client: ClientSessionGenerator,
+    sensor_entities: dict[str, er.RegistryEntry],
+    climate_entities: dict[str, er.RegistryEntry | dict[str, Any]],
 ) -> None:
     """Test disabling a entity."""
     data = {**sensor_entities, **climate_entities}
@@ -1760,14 +1812,14 @@ def mock_client_fixture():
         yield counter_client
 
 
-async def test_minimal_config(hass: HomeAssistant, mock_client) -> None:
+async def test_minimal_config(hass: HomeAssistant, mock_client: mock.MagicMock) -> None:
     """Test the minimal config and defaults of component."""
     config = {prometheus.DOMAIN: {}}
     assert await async_setup_component(hass, prometheus.DOMAIN, config)
     await hass.async_block_till_done()
 
 
-async def test_full_config(hass: HomeAssistant, mock_client) -> None:
+async def test_full_config(hass: HomeAssistant, mock_client: mock.MagicMock) -> None:
     """Test the full config of component."""
     config = {
         prometheus.DOMAIN: {
@@ -1792,14 +1844,14 @@ async def test_full_config(hass: HomeAssistant, mock_client) -> None:
     await hass.async_block_till_done()
 
 
-async def _setup(hass, filter_config):
+async def _setup(hass: HomeAssistant, filter_config):
     """Shared set up for filtering tests."""
     config = {prometheus.DOMAIN: {"filter": filter_config}}
     assert await async_setup_component(hass, prometheus.DOMAIN, config)
     await hass.async_block_till_done()
 
 
-async def test_allowlist(hass: HomeAssistant, mock_client) -> None:
+async def test_allowlist(hass: HomeAssistant, mock_client: mock.MagicMock) -> None:
     """Test an allowlist only config."""
     await _setup(
         hass,
@@ -1828,7 +1880,7 @@ async def test_allowlist(hass: HomeAssistant, mock_client) -> None:
         mock_client.labels.reset_mock()
 
 
-async def test_denylist(hass: HomeAssistant, mock_client) -> None:
+async def test_denylist(hass: HomeAssistant, mock_client: mock.MagicMock) -> None:
     """Test a denylist only config."""
     await _setup(
         hass,
@@ -1857,7 +1909,9 @@ async def test_denylist(hass: HomeAssistant, mock_client) -> None:
         mock_client.labels.reset_mock()
 
 
-async def test_filtered_denylist(hass: HomeAssistant, mock_client) -> None:
+async def test_filtered_denylist(
+    hass: HomeAssistant, mock_client: mock.MagicMock
+) -> None:
     """Test a denylist config with a filtering allowlist."""
     await _setup(
         hass,