From a7248d4574202bab7413b11ab54c95937ed5c026 Mon Sep 17 00:00:00 2001
From: Paulus Schoutsen <paulus@home-assistant.io>
Date: Thu, 27 Sep 2018 23:10:07 +0200
Subject: [PATCH] Handle exception handling websocket command (#16927)

* Handle exception handling websocket command

* lint

* Lint2
---
 homeassistant/components/websocket_api.py |  9 ++++++++-
 tests/components/test_websocket_api.py    | 19 ++++++++++++++++++-
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/homeassistant/components/websocket_api.py b/homeassistant/components/websocket_api.py
index 9bd4aac4b6a..4e7c186facc 100644
--- a/homeassistant/components/websocket_api.py
+++ b/homeassistant/components/websocket_api.py
@@ -40,6 +40,7 @@ ERR_ID_REUSE = 1
 ERR_INVALID_FORMAT = 2
 ERR_NOT_FOUND = 3
 ERR_UNKNOWN_COMMAND = 4
+ERR_UNKNOWN_ERROR = 5
 
 TYPE_AUTH = 'auth'
 TYPE_AUTH_INVALID = 'auth_invalid'
@@ -405,7 +406,13 @@ class ActiveConnection:
 
                 else:
                     handler, schema = handlers[msg['type']]
-                    handler(self.hass, self, schema(msg))
+                    try:
+                        handler(self.hass, self, schema(msg))
+                    except Exception:  # pylint: disable=broad-except
+                        _LOGGER.exception('Error handling message: %s', msg)
+                        self.to_write.put_nowait(error_message(
+                            cur_id, ERR_UNKNOWN_ERROR,
+                            'Unknown error.'))
 
                 last_id = cur_id
                 msg = await wsock.receive_json()
diff --git a/tests/components/test_websocket_api.py b/tests/components/test_websocket_api.py
index 199a9d804f8..cf74081adb1 100644
--- a/tests/components/test_websocket_api.py
+++ b/tests/components/test_websocket_api.py
@@ -1,6 +1,6 @@
 """Tests for the Home Assistant Websocket API."""
 import asyncio
-from unittest.mock import patch
+from unittest.mock import patch, Mock
 
 from aiohttp import WSMsgType
 from async_timeout import timeout
@@ -539,3 +539,20 @@ async def test_call_service_context_no_user(hass, aiohttp_client):
         assert call.service == 'test_service'
         assert call.data == {'hello': 'world'}
         assert call.context.user_id is None
+
+
+async def test_handler_failing(hass, websocket_client):
+    """Test a command that raises."""
+    hass.components.websocket_api.async_register_command(
+        'bla', Mock(side_effect=TypeError),
+        wapi.BASE_COMMAND_MESSAGE_SCHEMA.extend({'type': 'bla'}))
+    await websocket_client.send_json({
+        'id': 5,
+        'type': 'bla',
+    })
+
+    msg = await websocket_client.receive_json()
+    assert msg['id'] == 5
+    assert msg['type'] == wapi.TYPE_RESULT
+    assert not msg['success']
+    assert msg['error']['code'] == wapi.ERR_UNKNOWN_ERROR