From de27f24a4c9c6f1b3e225b4f3d30758a5120222f Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Thu, 13 Jun 2024 16:17:11 -0500 Subject: [PATCH] Use the existing api client for unifiprotect repairs if available (#119640) Co-authored-by: TheJulianJES --- homeassistant/components/unifiprotect/data.py | 4 +- .../components/unifiprotect/repairs.py | 42 +++++++------ tests/components/unifiprotect/test_repairs.py | 61 +++++++++++++++++++ 3 files changed, 88 insertions(+), 19 deletions(-) diff --git a/homeassistant/components/unifiprotect/data.py b/homeassistant/components/unifiprotect/data.py index 59e98cfb9a0..97f3a4129ae 100644 --- a/homeassistant/components/unifiprotect/data.py +++ b/homeassistant/components/unifiprotect/data.py @@ -376,7 +376,9 @@ def async_get_data_for_entry_id( hass: HomeAssistant, entry_id: str ) -> ProtectData | None: """Find the ProtectData instance for a config entry id.""" - if entry := hass.config_entries.async_get_entry(entry_id): + if (entry := hass.config_entries.async_get_entry(entry_id)) and hasattr( + entry, "runtime_data" + ): entry = cast(UFPConfigEntry, entry) return entry.runtime_data return None diff --git a/homeassistant/components/unifiprotect/repairs.py b/homeassistant/components/unifiprotect/repairs.py index 0e505f87391..020da0a03f6 100644 --- a/homeassistant/components/unifiprotect/repairs.py +++ b/homeassistant/components/unifiprotect/repairs.py @@ -11,11 +11,12 @@ import voluptuous as vol from homeassistant import data_entry_flow from homeassistant.components.repairs import ConfirmRepairFlow, RepairsFlow +from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant, callback from homeassistant.helpers import issue_registry as ir from .const import CONF_ALLOW_EA -from .data import UFPConfigEntry +from .data import UFPConfigEntry, async_get_data_for_entry_id from .utils import async_create_api_client @@ -219,29 +220,34 @@ class RTSPRepair(ProtectRepair): ) +@callback +def _async_get_or_create_api_client( + hass: HomeAssistant, entry: ConfigEntry +) -> ProtectApiClient: + """Get or create an API client.""" + if data := async_get_data_for_entry_id(hass, entry.entry_id): + return data.api + return async_create_api_client(hass, entry) + + async def async_create_fix_flow( hass: HomeAssistant, issue_id: str, data: dict[str, str | int | float | None] | None, ) -> RepairsFlow: """Create flow.""" - if data is not None and issue_id == "ea_channel_warning": - entry_id = cast(str, data["entry_id"]) - if (entry := hass.config_entries.async_get_entry(entry_id)) is not None: - api = async_create_api_client(hass, entry) + if ( + data is not None + and "entry_id" in data + and (entry := hass.config_entries.async_get_entry(cast(str, data["entry_id"]))) + ): + api = _async_get_or_create_api_client(hass, entry) + if issue_id == "ea_channel_warning": return EAConfirmRepair(api=api, entry=entry) - - elif data is not None and issue_id == "cloud_user": - entry_id = cast(str, data["entry_id"]) - if (entry := hass.config_entries.async_get_entry(entry_id)) is not None: - api = async_create_api_client(hass, entry) + if issue_id == "cloud_user": return CloudAccountRepair(api=api, entry=entry) - - elif data is not None and issue_id.startswith("rtsp_disabled_"): - entry_id = cast(str, data["entry_id"]) - camera_id = cast(str, data["camera_id"]) - if (entry := hass.config_entries.async_get_entry(entry_id)) is not None: - api = async_create_api_client(hass, entry) - return RTSPRepair(api=api, entry=entry, camera_id=camera_id) - + if issue_id.startswith("rtsp_disabled_"): + return RTSPRepair( + api=api, entry=entry, camera_id=cast(str, data["camera_id"]) + ) return ConfirmRepairFlow() diff --git a/tests/components/unifiprotect/test_repairs.py b/tests/components/unifiprotect/test_repairs.py index 7d76550f7c7..51ffd4d23cb 100644 --- a/tests/components/unifiprotect/test_repairs.py +++ b/tests/components/unifiprotect/test_repairs.py @@ -357,3 +357,64 @@ async def test_rtsp_writable_fix( ufp.api.update_device.assert_called_with( ModelType.CAMERA, doorbell.id, {"channels": channels} ) + + +async def test_rtsp_writable_fix_when_not_setup( + hass: HomeAssistant, + ufp: MockUFPFixture, + doorbell: Camera, + hass_client: ClientSessionGenerator, + hass_ws_client: WebSocketGenerator, +) -> None: + """Test RTSP disabled warning if the integration is no longer set up.""" + + for channel in doorbell.channels: + channel.is_rtsp_enabled = False + + await init_entry(hass, ufp, [doorbell]) + await async_process_repairs_platforms(hass) + ws_client = await hass_ws_client(hass) + client = await hass_client() + + new_doorbell = deepcopy(doorbell) + new_doorbell.channels[0].is_rtsp_enabled = True + ufp.api.get_camera = AsyncMock(side_effect=[doorbell, new_doorbell]) + ufp.api.update_device = AsyncMock() + issue_id = f"rtsp_disabled_{doorbell.id}" + + await ws_client.send_json({"id": 1, "type": "repairs/list_issues"}) + msg = await ws_client.receive_json() + + assert msg["success"] + assert len(msg["result"]["issues"]) > 0 + issue = None + for i in msg["result"]["issues"]: + if i["issue_id"] == issue_id: + issue = i + assert issue is not None + + # Unload the integration to ensure the fix flow still works + # if the integration is no longer set up + await hass.config_entries.async_unload(ufp.entry.entry_id) + await hass.async_block_till_done() + + url = RepairsFlowIndexView.url + resp = await client.post(url, json={"handler": DOMAIN, "issue_id": issue_id}) + assert resp.status == HTTPStatus.OK + data = await resp.json() + + flow_id = data["flow_id"] + assert data["step_id"] == "start" + + url = RepairsFlowResourceView.url.format(flow_id=flow_id) + resp = await client.post(url) + assert resp.status == HTTPStatus.OK + data = await resp.json() + + assert data["type"] == "create_entry" + + channels = doorbell.unifi_dict()["channels"] + channels[0]["isRtspEnabled"] = True + ufp.api.update_device.assert_called_with( + ModelType.CAMERA, doorbell.id, {"channels": channels} + )