From d66145358a4dae1c225e57f66531d372da15e394 Mon Sep 17 00:00:00 2001
From: Jan Bouwhuis <jbouwh@users.noreply.github.com>
Date: Wed, 3 Apr 2024 18:40:59 +0200
Subject: [PATCH] Correct imap services setup (#114760)

* Correct imap services setup

* Add config schema
---
 homeassistant/components/imap/__init__.py | 73 ++++++++++++-----------
 1 file changed, 37 insertions(+), 36 deletions(-)

diff --git a/homeassistant/components/imap/__init__.py b/homeassistant/components/imap/__init__.py
index 6c90889a7d6..22e32187255 100644
--- a/homeassistant/components/imap/__init__.py
+++ b/homeassistant/components/imap/__init__.py
@@ -19,6 +19,7 @@ from homeassistant.exceptions import (
     ServiceValidationError,
 )
 import homeassistant.helpers.config_validation as cv
+from homeassistant.helpers.typing import ConfigType
 
 from .const import CONF_ENABLE_PUSH, DOMAIN
 from .coordinator import (
@@ -37,6 +38,9 @@ CONF_TARGET_FOLDER = "target_folder"
 
 _LOGGER = logging.getLogger(__name__)
 
+
+CONFIG_SCHEMA = cv.config_entry_only_config_schema(DOMAIN)
+
 _SERVICE_UID_SCHEMA = vol.Schema(
     {
         vol.Required(CONF_ENTRY): cv.string,
@@ -95,30 +99,8 @@ def raise_on_error(response: Response, translation_key: str) -> None:
         )
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
-    """Set up imap from a config entry."""
-    try:
-        imap_client: IMAP4_SSL = await connect_to_server(dict(entry.data))
-    except InvalidAuth as err:
-        raise ConfigEntryAuthFailed from err
-    except InvalidFolder as err:
-        raise ConfigEntryError("Selected mailbox folder is invalid.") from err
-    except (TimeoutError, AioImapException) as err:
-        raise ConfigEntryNotReady from err
-
-    coordinator_class: type[
-        ImapPushDataUpdateCoordinator | ImapPollingDataUpdateCoordinator
-    ]
-    enable_push: bool = entry.data.get(CONF_ENABLE_PUSH, True)
-    if enable_push and imap_client.has_capability("IDLE"):
-        coordinator_class = ImapPushDataUpdateCoordinator
-    else:
-        coordinator_class = ImapPollingDataUpdateCoordinator
-
-    coordinator: ImapPushDataUpdateCoordinator | ImapPollingDataUpdateCoordinator = (
-        coordinator_class(hass, imap_client, entry)
-    )
-    await coordinator.async_config_entry_first_refresh()
+async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
+    """Set up imap services."""
 
     async def async_seen(call: ServiceCall) -> None:
         """Process mark as seen service call."""
@@ -141,8 +123,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         raise_on_error(response, "seen_failed")
         await client.close()
 
-    if not hass.services.has_service(DOMAIN, "seen"):
-        hass.services.async_register(DOMAIN, "seen", async_seen, SERVICE_SEEN_SCHEMA)
+    hass.services.async_register(DOMAIN, "seen", async_seen, SERVICE_SEEN_SCHEMA)
 
     async def async_move(call: ServiceCall) -> None:
         """Process move email service call."""
@@ -178,8 +159,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
             ) from exc
         await client.close()
 
-    if not hass.services.has_service(DOMAIN, "move"):
-        hass.services.async_register(DOMAIN, "move", async_move, SERVICE_MOVE_SCHEMA)
+    hass.services.async_register(DOMAIN, "move", async_move, SERVICE_MOVE_SCHEMA)
 
     async def async_delete(call: ServiceCall) -> None:
         """Process deleting email service call."""
@@ -206,10 +186,35 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
             ) from exc
         await client.close()
 
-    if not hass.services.has_service(DOMAIN, "delete"):
-        hass.services.async_register(
-            DOMAIN, "delete", async_delete, SERVICE_DELETE_SCHEMA
-        )
+    hass.services.async_register(DOMAIN, "delete", async_delete, SERVICE_DELETE_SCHEMA)
+
+    return True
+
+
+async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+    """Set up imap from a config entry."""
+    try:
+        imap_client: IMAP4_SSL = await connect_to_server(dict(entry.data))
+    except InvalidAuth as err:
+        raise ConfigEntryAuthFailed from err
+    except InvalidFolder as err:
+        raise ConfigEntryError("Selected mailbox folder is invalid.") from err
+    except (TimeoutError, AioImapException) as err:
+        raise ConfigEntryNotReady from err
+
+    coordinator_class: type[
+        ImapPushDataUpdateCoordinator | ImapPollingDataUpdateCoordinator
+    ]
+    enable_push: bool = entry.data.get(CONF_ENABLE_PUSH, True)
+    if enable_push and imap_client.has_capability("IDLE"):
+        coordinator_class = ImapPushDataUpdateCoordinator
+    else:
+        coordinator_class = ImapPollingDataUpdateCoordinator
+
+    coordinator: ImapPushDataUpdateCoordinator | ImapPollingDataUpdateCoordinator = (
+        coordinator_class(hass, imap_client, entry)
+    )
+    await coordinator.async_config_entry_first_refresh()
 
     hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
 
@@ -229,8 +234,4 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
             ImapPushDataUpdateCoordinator | ImapPollingDataUpdateCoordinator
         ) = hass.data[DOMAIN].pop(entry.entry_id)
         await coordinator.shutdown()
-    if not hass.data[DOMAIN]:
-        hass.services.async_remove(DOMAIN, "seen")
-        hass.services.async_remove(DOMAIN, "move")
-        hass.services.async_remove(DOMAIN, "delete")
     return unload_ok