From 673f542cdebfedc0c10bb239cc775c5d30a81a3f Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Fri, 16 Apr 2021 17:57:28 -1000
Subject: [PATCH] Do not wait for websocket response to be delivered before
 shutdown (#49323)

- Waiting was unreliable since the websocket response could take a few seconds to get delivered
- Alternate frontend fix is https://github.com/home-assistant/frontend/pull/8932
---
 .../components/homeassistant/__init__.py      | 26 ++-----------------
 tests/components/homeassistant/test_init.py   |  5 ----
 2 files changed, 2 insertions(+), 29 deletions(-)

diff --git a/homeassistant/components/homeassistant/__init__.py b/homeassistant/components/homeassistant/__init__.py
index 86be5862e7c..f80d3a0efb4 100644
--- a/homeassistant/components/homeassistant/__init__.py
+++ b/homeassistant/components/homeassistant/__init__.py
@@ -21,7 +21,6 @@ from homeassistant.const import (
 import homeassistant.core as ha
 from homeassistant.exceptions import HomeAssistantError, Unauthorized, UnknownUser
 from homeassistant.helpers import config_validation as cv, recorder
-from homeassistant.helpers.event import async_call_later
 from homeassistant.helpers.service import (
     async_extract_config_entry_ids,
     async_extract_referenced_entity_ids,
@@ -49,7 +48,6 @@ SCHEMA_RELOAD_CONFIG_ENTRY = vol.All(
 
 
 SHUTDOWN_SERVICES = (SERVICE_HOMEASSISTANT_STOP, SERVICE_HOMEASSISTANT_RESTART)
-WEBSOCKET_RECEIVE_DELAY = 1
 
 
 async def async_setup(hass: ha.HomeAssistant, config: dict) -> bool:
@@ -143,15 +141,7 @@ async def async_setup(hass: ha.HomeAssistant, config: dict) -> bool:
             )
 
         if call.service == SERVICE_HOMEASSISTANT_STOP:
-            # We delay the stop by WEBSOCKET_RECEIVE_DELAY to ensure the frontend
-            # can receive the response before the webserver shuts down
-            @ha.callback
-            def _async_stop(_):
-                # This must not be a tracked task otherwise
-                # the task itself will block stop
-                asyncio.create_task(hass.async_stop())
-
-            async_call_later(hass, WEBSOCKET_RECEIVE_DELAY, _async_stop)
+            asyncio.create_task(hass.async_stop())
             return
 
         errors = await conf_util.async_check_ha_config_file(hass)
@@ -172,19 +162,7 @@ async def async_setup(hass: ha.HomeAssistant, config: dict) -> bool:
             )
 
         if call.service == SERVICE_HOMEASSISTANT_RESTART:
-            # We delay the restart by WEBSOCKET_RECEIVE_DELAY to ensure the frontend
-            # can receive the response before the webserver shuts down
-            @ha.callback
-            def _async_stop_with_code(_):
-                # This must not be a tracked task otherwise
-                # the task itself will block restart
-                asyncio.create_task(hass.async_stop(RESTART_EXIT_CODE))
-
-            async_call_later(
-                hass,
-                WEBSOCKET_RECEIVE_DELAY,
-                _async_stop_with_code,
-            )
+            asyncio.create_task(hass.async_stop(RESTART_EXIT_CODE))
 
     async def async_handle_update_service(call):
         """Service handler for updating an entity."""
diff --git a/tests/components/homeassistant/test_init.py b/tests/components/homeassistant/test_init.py
index 451c226eb87..d12cc8d9a7b 100644
--- a/tests/components/homeassistant/test_init.py
+++ b/tests/components/homeassistant/test_init.py
@@ -1,7 +1,6 @@
 """The tests for Core components."""
 # pylint: disable=protected-access
 import asyncio
-from datetime import timedelta
 import unittest
 from unittest.mock import Mock, patch
 
@@ -34,12 +33,10 @@ import homeassistant.core as ha
 from homeassistant.exceptions import HomeAssistantError, Unauthorized
 from homeassistant.helpers import entity
 from homeassistant.setup import async_setup_component
-import homeassistant.util.dt as dt_util
 
 from tests.common import (
     MockConfigEntry,
     async_capture_events,
-    async_fire_time_changed,
     async_mock_service,
     get_test_home_assistant,
     mock_registry,
@@ -526,7 +523,6 @@ async def test_restart_homeassistant(hass):
             blocking=True,
         )
         assert mock_check.called
-        async_fire_time_changed(hass, dt_util.utcnow() + timedelta(seconds=2))
         await hass.async_block_till_done()
         assert mock_restart.called
 
@@ -545,6 +541,5 @@ async def test_stop_homeassistant(hass):
             blocking=True,
         )
         assert not mock_check.called
-        async_fire_time_changed(hass, dt_util.utcnow() + timedelta(seconds=2))
         await hass.async_block_till_done()
         assert mock_restart.called