Use HassKey in auth (#127573)

homeassistant/components/auth/__init__.py

@@ -159,6 +159,7 @@ from homeassistant.helpers.config_entry_oauth2_flow import OAuth2AuthorizeCallba
 from homeassistant.helpers.typing import ConfigType
 from homeassistant.loader import bind_hass
 from homeassistant.util import dt as dt_util
+from homeassistant.util.hass_dict import HassKey
 
 from . import indieauth, login_flow, mfa_setup_flow
 
@@ -166,7 +167,7 @@ DOMAIN = "auth"
 
 type StoreResultType = Callable[[str, Credentials], str]
 type RetrieveResultType = Callable[[str, str], Credentials | None]
-
+DATA_STORE: HassKey[StoreResultType] = HassKey(DOMAIN)
 CONFIG_SCHEMA = cv.empty_config_schema(DOMAIN)
 
 DELETE_CURRENT_TOKEN_DELAY = 2
@@ -177,14 +178,14 @@ def create_auth_code(
     hass: HomeAssistant, client_id: str, credential: Credentials
 ) -> str:
     """Create an authorization code to fetch tokens."""
-    return cast(StoreResultType, hass.data[DOMAIN])(client_id, credential)
+    return hass.data[DATA_STORE](client_id, credential)
 
 
 async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     """Component to allow users to login."""
     store_result, retrieve_result = _create_auth_code_store()
 
-    hass.data[DOMAIN] = store_result
+    hass.data[DATA_STORE] = store_result
 
     hass.http.register_view(TokenView(retrieve_result))
     hass.http.register_view(RevokeTokenView())

homeassistant/components/auth/mfa_setup_flow.py

@@ -12,6 +12,7 @@ from homeassistant import data_entry_flow
 from homeassistant.components import websocket_api
 from homeassistant.core import HomeAssistant, callback
 import homeassistant.helpers.config_validation as cv
+from homeassistant.util.hass_dict import HassKey
 
 WS_TYPE_SETUP_MFA = "auth/setup_mfa"
 SCHEMA_WS_SETUP_MFA = vol.All(
@@ -31,7 +32,7 @@ SCHEMA_WS_DEPOSE_MFA = websocket_api.BASE_COMMAND_MESSAGE_SCHEMA.extend(
     {vol.Required("type"): WS_TYPE_DEPOSE_MFA, vol.Required("mfa_module_id"): str}
 )
 
-DATA_SETUP_FLOW_MGR = "auth_mfa_setup_flow_manager"
+DATA_SETUP_FLOW_MGR: HassKey[MfaFlowManager] = HassKey("auth_mfa_setup_flow_manager")
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -89,7 +90,7 @@ def websocket_setup_mfa(
 
     async def async_setup_flow(msg: dict[str, Any]) -> None:
         """Return a setup flow for mfa auth module."""
-        flow_manager: MfaFlowManager = hass.data[DATA_SETUP_FLOW_MGR]
+        flow_manager = hass.data[DATA_SETUP_FLOW_MGR]
 
         if (flow_id := msg.get("flow_id")) is not None:
             result = await flow_manager.async_configure(flow_id, msg.get("user_input"))