Add type hints to integration tests (part 23) (#88235)

tests/components/trace/test_websocket_api.py

@@ -1,7 +1,7 @@
 """Test Trace websocket API."""
 import asyncio
 import json
-from typing import DefaultDict
+from typing import Any, DefaultDict
 from unittest.mock import patch
 
 import pytest
@@ -9,11 +9,12 @@ import pytest
 from homeassistant.bootstrap import async_setup_component
 from homeassistant.components.trace.const import DEFAULT_STORED_TRACES
 from homeassistant.const import EVENT_HOMEASSISTANT_STOP
-from homeassistant.core import Context, CoreState, callback
+from homeassistant.core import Context, CoreState, HomeAssistant, callback
 from homeassistant.helpers.typing import UNDEFINED
 from homeassistant.util.uuid import random_uuid_hex
 
 from tests.common import assert_lists_same, load_fixture
+from tests.typing import WebSocketGenerator
 
 
 def _find_run_id(traces, trace_type, item_id):
@@ -116,8 +117,8 @@ async def _assert_contexts(client, next_id, contexts, domain=None, item_id=None)
     ],
 )
 async def test_get_trace(
-    hass,
-    hass_storage,
+    hass: HomeAssistant,
+    hass_storage: dict[str, Any],
     hass_ws_client,
     domain,
     prefix,
@@ -125,8 +126,8 @@ async def test_get_trace(
     trigger,
     context_key,
     condition_results,
-    enable_custom_integrations,
-):
+    enable_custom_integrations: None,
+) -> None:
     """Test tracing a script or automation."""
     id = 1
 
@@ -420,7 +421,9 @@ async def test_get_trace(
 
 
 @pytest.mark.parametrize("domain", ["automation", "script"])
-async def test_restore_traces(hass, hass_storage, hass_ws_client, domain):
+async def test_restore_traces(
+    hass: HomeAssistant, hass_storage: dict[str, Any], hass_ws_client, domain
+) -> None:
     """Test restored traces."""
     hass.state = CoreState.not_running
     id = 1
@@ -488,7 +491,9 @@ async def test_restore_traces(hass, hass_storage, hass_ws_client, domain):
 
 
 @pytest.mark.parametrize("domain", ["automation", "script"])
-async def test_get_invalid_trace(hass, hass_ws_client, domain):
+async def test_get_invalid_trace(
+    hass: HomeAssistant, hass_ws_client: WebSocketGenerator, domain
+) -> None:
     """Test getting a non-existing trace."""
     assert await async_setup_component(hass, domain, {domain: {}})
     client = await hass_ws_client()
@@ -510,7 +515,9 @@ async def test_get_invalid_trace(hass, hass_ws_client, domain):
     ("domain", "stored_traces"),
     [("automation", None), ("automation", 10), ("script", None), ("script", 10)],
 )
-async def test_trace_overflow(hass, hass_ws_client, domain, stored_traces):
+async def test_trace_overflow(
+    hass: HomeAssistant, hass_ws_client: WebSocketGenerator, domain, stored_traces
+) -> None:
     """Test the number of stored traces per script or automation is limited."""
     id = 1
 
@@ -583,8 +590,12 @@ async def test_trace_overflow(hass, hass_ws_client, domain, stored_traces):
     ("domain", "num_restored_moon_traces"), [("automation", 3), ("script", 1)]
 )
 async def test_restore_traces_overflow(
-    hass, hass_storage, hass_ws_client, domain, num_restored_moon_traces
-):
+    hass: HomeAssistant,
+    hass_storage: dict[str, Any],
+    hass_ws_client,
+    domain,
+    num_restored_moon_traces,
+) -> None:
     """Test restored traces are evicted first."""
     hass.state = CoreState.not_running
     id = 1
@@ -659,13 +670,13 @@ async def test_restore_traces_overflow(
     [("automation", 3, "e2c97432afe9b8a42d7983588ed5e6ef"), ("script", 1, "")],
 )
 async def test_restore_traces_late_overflow(
-    hass,
-    hass_storage,
+    hass: HomeAssistant,
+    hass_storage: dict[str, Any],
     hass_ws_client,
     domain,
     num_restored_moon_traces,
     restored_run_id,
-):
+) -> None:
     """Test restored traces are evicted first."""
     hass.state = CoreState.not_running
     id = 1
@@ -725,7 +736,9 @@ async def test_restore_traces_late_overflow(
 
 
 @pytest.mark.parametrize("domain", ["automation", "script"])
-async def test_trace_no_traces(hass, hass_ws_client, domain):
+async def test_trace_no_traces(
+    hass: HomeAssistant, hass_ws_client: WebSocketGenerator, domain
+) -> None:
     """Test the storing traces for a script or automation can be disabled."""
     id = 1
 
@@ -784,8 +797,14 @@ async def test_trace_no_traces(hass, hass_ws_client, domain):
     ],
 )
 async def test_list_traces(
-    hass, hass_ws_client, domain, prefix, trigger, last_step, script_execution
-):
+    hass: HomeAssistant,
+    hass_ws_client: WebSocketGenerator,
+    domain,
+    prefix,
+    trigger,
+    last_step,
+    script_execution,
+) -> None:
     """Test listing script and automation traces."""
     id = 1
 
@@ -912,7 +931,13 @@ async def test_list_traces(
     ("domain", "prefix", "extra_trace_keys"),
     [("automation", "action", {"trigger/0"}), ("script", "sequence", set())],
 )
-async def test_nested_traces(hass, hass_ws_client, domain, prefix, extra_trace_keys):
+async def test_nested_traces(
+    hass: HomeAssistant,
+    hass_ws_client: WebSocketGenerator,
+    domain,
+    prefix,
+    extra_trace_keys,
+) -> None:
     """Test nested automation and script traces."""
     id = 1
 
@@ -970,7 +995,9 @@ async def test_nested_traces(hass, hass_ws_client, domain, prefix, extra_trace_k
 @pytest.mark.parametrize(
     ("domain", "prefix"), [("automation", "action"), ("script", "sequence")]
 )
-async def test_breakpoints(hass, hass_ws_client, domain, prefix):
+async def test_breakpoints(
+    hass: HomeAssistant, hass_ws_client: WebSocketGenerator, domain, prefix
+) -> None:
     """Test script and automation breakpoints."""
     id = 1
 
@@ -1139,7 +1166,9 @@ async def test_breakpoints(hass, hass_ws_client, domain, prefix):
 @pytest.mark.parametrize(
     ("domain", "prefix"), [("automation", "action"), ("script", "sequence")]
 )
-async def test_breakpoints_2(hass, hass_ws_client, domain, prefix):
+async def test_breakpoints_2(
+    hass: HomeAssistant, hass_ws_client: WebSocketGenerator, domain, prefix
+) -> None:
     """Test execution resumes and breakpoints are removed after subscription removed."""
     id = 1
 
@@ -1242,7 +1271,9 @@ async def test_breakpoints_2(hass, hass_ws_client, domain, prefix):
 @pytest.mark.parametrize(
     ("domain", "prefix"), [("automation", "action"), ("script", "sequence")]
 )
-async def test_breakpoints_3(hass, hass_ws_client, domain, prefix):
+async def test_breakpoints_3(
+    hass: HomeAssistant, hass_ws_client: WebSocketGenerator, domain, prefix
+) -> None:
     """Test breakpoints can be cleared."""
     id = 1
 
@@ -1393,8 +1424,12 @@ async def test_breakpoints_3(hass, hass_ws_client, domain, prefix):
     ],
 )
 async def test_script_mode(
-    hass, hass_ws_client, script_mode, max_runs, script_execution
-):
+    hass: HomeAssistant,
+    hass_ws_client: WebSocketGenerator,
+    script_mode,
+    max_runs,
+    script_execution,
+) -> None:
     """Test overlapping runs with max_runs > 1."""
     id = 1
 
@@ -1457,7 +1492,12 @@ async def test_script_mode(
     ("script_mode", "script_execution"),
     [("restart", "cancelled"), ("parallel", "finished")],
 )
-async def test_script_mode_2(hass, hass_ws_client, script_mode, script_execution):
+async def test_script_mode_2(
+    hass: HomeAssistant,
+    hass_ws_client: WebSocketGenerator,
+    script_mode,
+    script_execution,
+) -> None:
     """Test overlapping runs with max_runs > 1."""
     id = 1
 
@@ -1528,8 +1568,10 @@ async def test_script_mode_2(hass, hass_ws_client, script_mode, script_execution
 
 
 async def test_trace_blueprint_automation(
-    hass, hass_ws_client, enable_custom_integrations
-):
+    hass: HomeAssistant,
+    hass_ws_client: WebSocketGenerator,
+    enable_custom_integrations: None,
+) -> None:
     """Test trace of blueprint automation."""
     id = 1