From 40ce25800ca4649ebb3047f57bd91e4f1b991213 Mon Sep 17 00:00:00 2001
From: Erik Montnemery <erik@montnemery.com>
Date: Mon, 22 Mar 2021 06:12:56 +0100
Subject: [PATCH] Test that homeassistant stop and restart do not block WS
 (#48081)

---
 .../components/websocket_api/test_commands.py | 80 +++++++++++++++++++
 1 file changed, 80 insertions(+)

diff --git a/tests/components/websocket_api/test_commands.py b/tests/components/websocket_api/test_commands.py
index a83f9509d01..b9e1e149dd1 100644
--- a/tests/components/websocket_api/test_commands.py
+++ b/tests/components/websocket_api/test_commands.py
@@ -1,5 +1,8 @@
 """Tests for WebSocket API commands."""
+from unittest.mock import ANY, patch
+
 from async_timeout import timeout
+import pytest
 import voluptuous as vol
 
 from homeassistant.components.websocket_api import const
@@ -47,6 +50,83 @@ async def test_call_service(hass, websocket_client):
     assert call.context.as_dict() == msg["result"]["context"]
 
 
+@pytest.mark.parametrize("command", ("call_service", "call_service_action"))
+async def test_call_service_blocking(hass, websocket_client, command):
+    """Test call service commands block, except for homeassistant restart / stop."""
+    with patch(
+        "homeassistant.core.ServiceRegistry.async_call", autospec=True
+    ) as mock_call:
+        await websocket_client.send_json(
+            {
+                "id": 5,
+                "type": "call_service",
+                "domain": "domain_test",
+                "service": "test_service",
+                "service_data": {"hello": "world"},
+            },
+        )
+        msg = await websocket_client.receive_json()
+
+    assert msg["id"] == 5
+    assert msg["type"] == const.TYPE_RESULT
+    assert msg["success"]
+    mock_call.assert_called_once_with(
+        ANY,
+        "domain_test",
+        "test_service",
+        {"hello": "world"},
+        blocking=True,
+        context=ANY,
+        target=ANY,
+    )
+
+    with patch(
+        "homeassistant.core.ServiceRegistry.async_call", autospec=True
+    ) as mock_call:
+        await websocket_client.send_json(
+            {
+                "id": 6,
+                "type": "call_service",
+                "domain": "homeassistant",
+                "service": "test_service",
+            },
+        )
+        msg = await websocket_client.receive_json()
+
+    assert msg["id"] == 6
+    assert msg["type"] == const.TYPE_RESULT
+    assert msg["success"]
+    mock_call.assert_called_once_with(
+        ANY,
+        "homeassistant",
+        "test_service",
+        ANY,
+        blocking=True,
+        context=ANY,
+        target=ANY,
+    )
+
+    with patch(
+        "homeassistant.core.ServiceRegistry.async_call", autospec=True
+    ) as mock_call:
+        await websocket_client.send_json(
+            {
+                "id": 7,
+                "type": "call_service",
+                "domain": "homeassistant",
+                "service": "restart",
+            },
+        )
+        msg = await websocket_client.receive_json()
+
+    assert msg["id"] == 7
+    assert msg["type"] == const.TYPE_RESULT
+    assert msg["success"]
+    mock_call.assert_called_once_with(
+        ANY, "homeassistant", "restart", ANY, blocking=False, context=ANY, target=ANY
+    )
+
+
 async def test_call_service_target(hass, websocket_client):
     """Test call service command with target."""
     calls = async_mock_service(hass, "domain_test", "test_service")