Skip to content

Commit

Permalink
Use new style for built-in ws commmands (home-assistant#21694)
Browse files Browse the repository at this point in the history
* Use new style for built-in ws commmands

* Lint
  • Loading branch information
balloob authored Mar 6, 2019
1 parent c9b1734 commit fc1ee9b
Show file tree
Hide file tree
Showing 4 changed files with 55 additions and 82 deletions.
96 changes: 35 additions & 61 deletions homeassistant/components/websocket_api/commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,77 +10,25 @@

from . import const, decorators, messages

TYPE_CALL_SERVICE = 'call_service'
TYPE_EVENT = 'event'
TYPE_GET_CONFIG = 'get_config'
TYPE_GET_SERVICES = 'get_services'
TYPE_GET_STATES = 'get_states'
TYPE_PING = 'ping'
TYPE_PONG = 'pong'
TYPE_SUBSCRIBE_EVENTS = 'subscribe_events'
TYPE_UNSUBSCRIBE_EVENTS = 'unsubscribe_events'


@callback
def async_register_commands(hass):
"""Register commands."""
async_reg = hass.components.websocket_api.async_register_command
async_reg(TYPE_SUBSCRIBE_EVENTS, handle_subscribe_events,
SCHEMA_SUBSCRIBE_EVENTS)
async_reg(TYPE_UNSUBSCRIBE_EVENTS, handle_unsubscribe_events,
SCHEMA_UNSUBSCRIBE_EVENTS)
async_reg(TYPE_CALL_SERVICE, handle_call_service, SCHEMA_CALL_SERVICE)
async_reg(TYPE_GET_STATES, handle_get_states, SCHEMA_GET_STATES)
async_reg(TYPE_GET_SERVICES, handle_get_services, SCHEMA_GET_SERVICES)
async_reg(TYPE_GET_CONFIG, handle_get_config, SCHEMA_GET_CONFIG)
async_reg(TYPE_PING, handle_ping, SCHEMA_PING)


SCHEMA_SUBSCRIBE_EVENTS = messages.BASE_COMMAND_MESSAGE_SCHEMA.extend({
vol.Required('type'): TYPE_SUBSCRIBE_EVENTS,
vol.Optional('event_type', default=MATCH_ALL): str,
})


SCHEMA_UNSUBSCRIBE_EVENTS = messages.BASE_COMMAND_MESSAGE_SCHEMA.extend({
vol.Required('type'): TYPE_UNSUBSCRIBE_EVENTS,
vol.Required('subscription'): cv.positive_int,
})


SCHEMA_CALL_SERVICE = messages.BASE_COMMAND_MESSAGE_SCHEMA.extend({
vol.Required('type'): TYPE_CALL_SERVICE,
vol.Required('domain'): str,
vol.Required('service'): str,
vol.Optional('service_data'): dict
})


SCHEMA_GET_STATES = messages.BASE_COMMAND_MESSAGE_SCHEMA.extend({
vol.Required('type'): TYPE_GET_STATES,
})


SCHEMA_GET_SERVICES = messages.BASE_COMMAND_MESSAGE_SCHEMA.extend({
vol.Required('type'): TYPE_GET_SERVICES,
})


SCHEMA_GET_CONFIG = messages.BASE_COMMAND_MESSAGE_SCHEMA.extend({
vol.Required('type'): TYPE_GET_CONFIG,
})


SCHEMA_PING = messages.BASE_COMMAND_MESSAGE_SCHEMA.extend({
vol.Required('type'): TYPE_PING,
})
async_reg(handle_subscribe_events)
async_reg(handle_unsubscribe_events)
async_reg(handle_call_service)
async_reg(handle_get_states)
async_reg(handle_get_services)
async_reg(handle_get_config)
async_reg(handle_ping)


def event_message(iden, event):
"""Return an event message."""
return {
'id': iden,
'type': TYPE_EVENT,
'type': 'event',
'event': event.as_dict(),
}

Expand All @@ -89,11 +37,15 @@ def pong_message(iden):
"""Return a pong message."""
return {
'id': iden,
'type': TYPE_PONG,
'type': 'pong',
}


@callback
@decorators.websocket_command({
vol.Required('type'): 'subscribe_events',
vol.Optional('event_type', default=MATCH_ALL): str,
})
def handle_subscribe_events(hass, connection, msg):
"""Handle subscribe events command.
Expand All @@ -116,6 +68,10 @@ async def forward_events(event):


@callback
@decorators.websocket_command({
vol.Required('type'): 'unsubscribe_events',
vol.Required('subscription'): cv.positive_int,
})
def handle_unsubscribe_events(hass, connection, msg):
"""Handle unsubscribe events command.
Expand All @@ -132,6 +88,12 @@ def handle_unsubscribe_events(hass, connection, msg):


@decorators.async_response
@decorators.websocket_command({
vol.Required('type'): 'call_service',
vol.Required('domain'): str,
vol.Required('service'): str,
vol.Optional('service_data'): dict
})
async def handle_call_service(hass, connection, msg):
"""Handle call service command.
Expand Down Expand Up @@ -161,6 +123,9 @@ async def handle_call_service(hass, connection, msg):


@callback
@decorators.websocket_command({
vol.Required('type'): 'get_states',
})
def handle_get_states(hass, connection, msg):
"""Handle get states command.
Expand All @@ -177,6 +142,9 @@ def handle_get_states(hass, connection, msg):


@decorators.async_response
@decorators.websocket_command({
vol.Required('type'): 'get_services',
})
async def handle_get_services(hass, connection, msg):
"""Handle get services command.
Expand All @@ -188,6 +156,9 @@ async def handle_get_services(hass, connection, msg):


@callback
@decorators.websocket_command({
vol.Required('type'): 'get_config',
})
def handle_get_config(hass, connection, msg):
"""Handle get config command.
Expand All @@ -198,6 +169,9 @@ def handle_get_config(hass, connection, msg):


@callback
@decorators.websocket_command({
vol.Required('type'): 'ping',
})
def handle_ping(hass, connection, msg):
"""Handle ping command.
Expand Down
3 changes: 1 addition & 2 deletions tests/components/websocket_api/test_auth.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@
from homeassistant.components.websocket_api.auth import (
TYPE_AUTH, TYPE_AUTH_INVALID, TYPE_AUTH_OK, TYPE_AUTH_REQUIRED)

from homeassistant.components.websocket_api import commands
from homeassistant.setup import async_setup_component

from tests.common import mock_coro
Expand Down Expand Up @@ -45,7 +44,7 @@ async def test_auth_via_msg_incorrect_pass(no_auth_websocket_client):
async def test_pre_auth_only_auth_allowed(no_auth_websocket_client):
"""Verify that before authentication, only auth messages are allowed."""
await no_auth_websocket_client.send_json({
'type': commands.TYPE_CALL_SERVICE,
'type': 'call_service',
'domain': 'domain_test',
'service': 'test_service',
'service_data': {
Expand Down
34 changes: 17 additions & 17 deletions tests/components/websocket_api/test_commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
from homeassistant.components.websocket_api.auth import (
TYPE_AUTH, TYPE_AUTH_OK, TYPE_AUTH_REQUIRED
)
from homeassistant.components.websocket_api import const, commands
from homeassistant.components.websocket_api import const
from homeassistant.exceptions import HomeAssistantError
from homeassistant.setup import async_setup_component

Expand All @@ -27,7 +27,7 @@ def service_call(call):

await websocket_client.send_json({
'id': 5,
'type': commands.TYPE_CALL_SERVICE,
'type': 'call_service',
'domain': 'domain_test',
'service': 'test_service',
'service_data': {
Expand All @@ -52,7 +52,7 @@ async def test_call_service_not_found(hass, websocket_client):
"""Test call service command."""
await websocket_client.send_json({
'id': 5,
'type': commands.TYPE_CALL_SERVICE,
'type': 'call_service',
'domain': 'domain_test',
'service': 'test_service',
'service_data': {
Expand Down Expand Up @@ -83,7 +83,7 @@ async def unknown_error_call(_):

await websocket_client.send_json({
'id': 5,
'type': commands.TYPE_CALL_SERVICE,
'type': 'call_service',
'domain': 'domain_test',
'service': 'ha_error',
})
Expand All @@ -98,7 +98,7 @@ async def unknown_error_call(_):

await websocket_client.send_json({
'id': 6,
'type': commands.TYPE_CALL_SERVICE,
'type': 'call_service',
'domain': 'domain_test',
'service': 'unknown_error',
})
Expand All @@ -118,7 +118,7 @@ async def test_subscribe_unsubscribe_events(hass, websocket_client):

await websocket_client.send_json({
'id': 5,
'type': commands.TYPE_SUBSCRIBE_EVENTS,
'type': 'subscribe_events',
'event_type': 'test_event'
})

Expand All @@ -138,7 +138,7 @@ async def test_subscribe_unsubscribe_events(hass, websocket_client):
msg = await websocket_client.receive_json()

assert msg['id'] == 5
assert msg['type'] == commands.TYPE_EVENT
assert msg['type'] == 'event'
event = msg['event']

assert event['event_type'] == 'test_event'
Expand All @@ -147,7 +147,7 @@ async def test_subscribe_unsubscribe_events(hass, websocket_client):

await websocket_client.send_json({
'id': 6,
'type': commands.TYPE_UNSUBSCRIBE_EVENTS,
'type': 'unsubscribe_events',
'subscription': 5
})

Expand All @@ -167,7 +167,7 @@ async def test_get_states(hass, websocket_client):

await websocket_client.send_json({
'id': 5,
'type': commands.TYPE_GET_STATES,
'type': 'get_states',
})

msg = await websocket_client.receive_json()
Expand All @@ -189,7 +189,7 @@ async def test_get_services(hass, websocket_client):
"""Test get_services command."""
await websocket_client.send_json({
'id': 5,
'type': commands.TYPE_GET_SERVICES,
'type': 'get_services',
})

msg = await websocket_client.receive_json()
Expand All @@ -203,7 +203,7 @@ async def test_get_config(hass, websocket_client):
"""Test get_config command."""
await websocket_client.send_json({
'id': 5,
'type': commands.TYPE_GET_CONFIG,
'type': 'get_config',
})

msg = await websocket_client.receive_json()
Expand All @@ -224,12 +224,12 @@ async def test_ping(websocket_client):
"""Test get_panels command."""
await websocket_client.send_json({
'id': 5,
'type': commands.TYPE_PING,
'type': 'ping',
})

msg = await websocket_client.receive_json()
assert msg['id'] == 5
assert msg['type'] == commands.TYPE_PONG
assert msg['type'] == 'pong'


async def test_call_service_context_with_user(hass, aiohttp_client,
Expand Down Expand Up @@ -258,7 +258,7 @@ async def test_call_service_context_with_user(hass, aiohttp_client,

await ws.send_json({
'id': 5,
'type': commands.TYPE_CALL_SERVICE,
'type': 'call_service',
'domain': 'domain_test',
'service': 'test_service',
'service_data': {
Expand All @@ -285,7 +285,7 @@ async def test_subscribe_requires_admin(websocket_client, hass_admin_user):
hass_admin_user.groups = []
await websocket_client.send_json({
'id': 5,
'type': commands.TYPE_SUBSCRIBE_EVENTS,
'type': 'subscribe_events',
'event_type': 'test_event'
})

Expand All @@ -307,7 +307,7 @@ async def test_states_filters_visible(hass, hass_admin_user, websocket_client):
hass.states.async_set('test.not_visible_entity', 'invisible')
await websocket_client.send_json({
'id': 5,
'type': commands.TYPE_GET_STATES,
'type': 'get_states',
})

msg = await websocket_client.receive_json()
Expand All @@ -327,7 +327,7 @@ async def test_get_states_not_allows_nan(hass, websocket_client):

await websocket_client.send_json({
'id': 5,
'type': commands.TYPE_GET_STATES,
'type': 'get_states',
})

msg = await websocket_client.receive_json()
Expand Down
4 changes: 2 additions & 2 deletions tests/components/websocket_api/test_init.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
from aiohttp import WSMsgType
import pytest

from homeassistant.components.websocket_api import const, commands, messages
from homeassistant.components.websocket_api import const, messages


@pytest.fixture
Expand Down Expand Up @@ -56,7 +56,7 @@ def test_pending_msg_overflow(hass, mock_low_queue, websocket_client):
for idx in range(10):
yield from websocket_client.send_json({
'id': idx + 1,
'type': commands.TYPE_PING,
'type': 'ping',
})
msg = yield from websocket_client.receive()
assert msg.type == WSMsgType.close
Expand Down

0 comments on commit fc1ee9b

Please sign in to comment.