diff --git a/homeassistant/components/mqtt_statestream.py b/homeassistant/components/mqtt_statestream.py index d24361637e9..fa1da879110 100644 --- a/homeassistant/components/mqtt_statestream.py +++ b/homeassistant/components/mqtt_statestream.py @@ -9,9 +9,11 @@ import json import voluptuous as vol -from homeassistant.const import MATCH_ALL +from homeassistant.const import (CONF_DOMAINS, CONF_ENTITIES, CONF_EXCLUDE, + CONF_INCLUDE, MATCH_ALL) from homeassistant.core import callback from homeassistant.components.mqtt import valid_publish_topic +from homeassistant.helpers.entityfilter import generate_filter from homeassistant.helpers.event import async_track_state_change from homeassistant.remote import JSONEncoder import homeassistant.helpers.config_validation as cv @@ -23,7 +25,7 @@ DEPENDENCIES = ['mqtt'] DOMAIN = 'mqtt_statestream' CONFIG_SCHEMA = vol.Schema({ - DOMAIN: vol.Schema({ + DOMAIN: cv.FILTER_SCHEMA.extend({ vol.Required(CONF_BASE_TOPIC): valid_publish_topic, vol.Optional(CONF_PUBLISH_ATTRIBUTES, default=False): cv.boolean, vol.Optional(CONF_PUBLISH_TIMESTAMPS, default=False): cv.boolean @@ -36,8 +38,14 @@ def async_setup(hass, config): """Set up the MQTT state feed.""" conf = config.get(DOMAIN, {}) base_topic = conf.get(CONF_BASE_TOPIC) + pub_include = conf.get(CONF_INCLUDE, {}) + pub_exclude = conf.get(CONF_EXCLUDE, {}) publish_attributes = conf.get(CONF_PUBLISH_ATTRIBUTES) publish_timestamps = conf.get(CONF_PUBLISH_TIMESTAMPS) + publish_filter = generate_filter(pub_include.get(CONF_DOMAINS, []), + pub_include.get(CONF_ENTITIES, []), + pub_exclude.get(CONF_DOMAINS, []), + pub_exclude.get(CONF_ENTITIES, [])) if not base_topic.endswith('/'): base_topic = base_topic + '/' @@ -45,6 +53,10 @@ def async_setup(hass, config): def _state_publisher(entity_id, old_state, new_state): if new_state is None: return + + if not publish_filter(entity_id): + return + payload = new_state.state mybase = base_topic + entity_id.replace('.', '/') + '/' diff --git a/tests/components/test_mqtt_statestream.py b/tests/components/test_mqtt_statestream.py index cc1ea277a34..76d8e48d03a 100644 --- a/tests/components/test_mqtt_statestream.py +++ b/tests/components/test_mqtt_statestream.py @@ -25,7 +25,8 @@ class TestMqttStateStream(object): self.hass.stop() def add_statestream(self, base_topic=None, publish_attributes=None, - publish_timestamps=None): + publish_timestamps=None, publish_include=None, + publish_exclude=None): """Add a mqtt_statestream component.""" config = {} if base_topic: @@ -34,7 +35,10 @@ class TestMqttStateStream(object): config['publish_attributes'] = publish_attributes if publish_timestamps: config['publish_timestamps'] = publish_timestamps - print("Publishing timestamps") + if publish_include: + config['include'] = publish_include + if publish_exclude: + config['exclude'] = publish_exclude return setup_component(self.hass, statestream.DOMAIN, { statestream.DOMAIN: config}) @@ -152,3 +156,237 @@ class TestMqttStateStream(object): mock_pub.assert_has_calls(calls, any_order=True) assert mock_pub.called + + @patch('homeassistant.components.mqtt.async_publish') + @patch('homeassistant.core.dt_util.utcnow') + def test_state_changed_event_include_domain(self, mock_utcnow, mock_pub): + """"Test that filtering on included domain works as expected.""" + base_topic = 'pub' + + incl = { + 'domains': ['fake'] + } + excl = {} + + # Add the statestream component for publishing state updates + # Set the filter to allow fake.* items + assert self.add_statestream(base_topic=base_topic, + publish_include=incl, + publish_exclude=excl) + self.hass.block_till_done() + + # Reset the mock because it will have already gotten calls for the + # mqtt_statestream state change on initialization, etc. + mock_pub.reset_mock() + + # Set a state of an entity + mock_state_change_event(self.hass, State('fake.entity', 'on')) + self.hass.block_till_done() + + # Make sure 'on' was published to pub/fake/entity/state + mock_pub.assert_called_with(self.hass, 'pub/fake/entity/state', 'on', + 1, True) + assert mock_pub.called + + mock_pub.reset_mock() + # Set a state of an entity that shouldn't be included + mock_state_change_event(self.hass, State('fake2.entity', 'on')) + self.hass.block_till_done() + + assert not mock_pub.called + + @patch('homeassistant.components.mqtt.async_publish') + @patch('homeassistant.core.dt_util.utcnow') + def test_state_changed_event_include_entity(self, mock_utcnow, mock_pub): + """"Test that filtering on included entity works as expected.""" + base_topic = 'pub' + + incl = { + 'entities': ['fake.entity'] + } + excl = {} + + # Add the statestream component for publishing state updates + # Set the filter to allow fake.* items + assert self.add_statestream(base_topic=base_topic, + publish_include=incl, + publish_exclude=excl) + self.hass.block_till_done() + + # Reset the mock because it will have already gotten calls for the + # mqtt_statestream state change on initialization, etc. + mock_pub.reset_mock() + + # Set a state of an entity + mock_state_change_event(self.hass, State('fake.entity', 'on')) + self.hass.block_till_done() + + # Make sure 'on' was published to pub/fake/entity/state + mock_pub.assert_called_with(self.hass, 'pub/fake/entity/state', 'on', + 1, True) + assert mock_pub.called + + mock_pub.reset_mock() + # Set a state of an entity that shouldn't be included + mock_state_change_event(self.hass, State('fake.entity2', 'on')) + self.hass.block_till_done() + + assert not mock_pub.called + + @patch('homeassistant.components.mqtt.async_publish') + @patch('homeassistant.core.dt_util.utcnow') + def test_state_changed_event_exclude_domain(self, mock_utcnow, mock_pub): + """"Test that filtering on excluded domain works as expected.""" + base_topic = 'pub' + + incl = {} + excl = { + 'domains': ['fake2'] + } + + # Add the statestream component for publishing state updates + # Set the filter to allow fake.* items + assert self.add_statestream(base_topic=base_topic, + publish_include=incl, + publish_exclude=excl) + self.hass.block_till_done() + + # Reset the mock because it will have already gotten calls for the + # mqtt_statestream state change on initialization, etc. + mock_pub.reset_mock() + + # Set a state of an entity + mock_state_change_event(self.hass, State('fake.entity', 'on')) + self.hass.block_till_done() + + # Make sure 'on' was published to pub/fake/entity/state + mock_pub.assert_called_with(self.hass, 'pub/fake/entity/state', 'on', + 1, True) + assert mock_pub.called + + mock_pub.reset_mock() + # Set a state of an entity that shouldn't be included + mock_state_change_event(self.hass, State('fake2.entity', 'on')) + self.hass.block_till_done() + + assert not mock_pub.called + + @patch('homeassistant.components.mqtt.async_publish') + @patch('homeassistant.core.dt_util.utcnow') + def test_state_changed_event_exclude_entity(self, mock_utcnow, mock_pub): + """"Test that filtering on excluded entity works as expected.""" + base_topic = 'pub' + + incl = {} + excl = { + 'entities': ['fake.entity2'] + } + + # Add the statestream component for publishing state updates + # Set the filter to allow fake.* items + assert self.add_statestream(base_topic=base_topic, + publish_include=incl, + publish_exclude=excl) + self.hass.block_till_done() + + # Reset the mock because it will have already gotten calls for the + # mqtt_statestream state change on initialization, etc. + mock_pub.reset_mock() + + # Set a state of an entity + mock_state_change_event(self.hass, State('fake.entity', 'on')) + self.hass.block_till_done() + + # Make sure 'on' was published to pub/fake/entity/state + mock_pub.assert_called_with(self.hass, 'pub/fake/entity/state', 'on', + 1, True) + assert mock_pub.called + + mock_pub.reset_mock() + # Set a state of an entity that shouldn't be included + mock_state_change_event(self.hass, State('fake.entity2', 'on')) + self.hass.block_till_done() + + assert not mock_pub.called + + @patch('homeassistant.components.mqtt.async_publish') + @patch('homeassistant.core.dt_util.utcnow') + def test_state_changed_event_exclude_domain_include_entity( + self, mock_utcnow, mock_pub): + """"Test filtering with excluded domain and included entity.""" + base_topic = 'pub' + + incl = { + 'entities': ['fake.entity'] + } + excl = { + 'domains': ['fake'] + } + + # Add the statestream component for publishing state updates + # Set the filter to allow fake.* items + assert self.add_statestream(base_topic=base_topic, + publish_include=incl, + publish_exclude=excl) + self.hass.block_till_done() + + # Reset the mock because it will have already gotten calls for the + # mqtt_statestream state change on initialization, etc. + mock_pub.reset_mock() + + # Set a state of an entity + mock_state_change_event(self.hass, State('fake.entity', 'on')) + self.hass.block_till_done() + + # Make sure 'on' was published to pub/fake/entity/state + mock_pub.assert_called_with(self.hass, 'pub/fake/entity/state', 'on', + 1, True) + assert mock_pub.called + + mock_pub.reset_mock() + # Set a state of an entity that shouldn't be included + mock_state_change_event(self.hass, State('fake.entity2', 'on')) + self.hass.block_till_done() + + assert not mock_pub.called + + @patch('homeassistant.components.mqtt.async_publish') + @patch('homeassistant.core.dt_util.utcnow') + def test_state_changed_event_include_domain_exclude_entity( + self, mock_utcnow, mock_pub): + """"Test filtering with included domain and excluded entity.""" + base_topic = 'pub' + + incl = { + 'domains': ['fake'] + } + excl = { + 'entities': ['fake.entity2'] + } + + # Add the statestream component for publishing state updates + # Set the filter to allow fake.* items + assert self.add_statestream(base_topic=base_topic, + publish_include=incl, + publish_exclude=excl) + self.hass.block_till_done() + + # Reset the mock because it will have already gotten calls for the + # mqtt_statestream state change on initialization, etc. + mock_pub.reset_mock() + + # Set a state of an entity + mock_state_change_event(self.hass, State('fake.entity', 'on')) + self.hass.block_till_done() + + # Make sure 'on' was published to pub/fake/entity/state + mock_pub.assert_called_with(self.hass, 'pub/fake/entity/state', 'on', + 1, True) + assert mock_pub.called + + mock_pub.reset_mock() + # Set a state of an entity that shouldn't be included + mock_state_change_event(self.hass, State('fake.entity2', 'on')) + self.hass.block_till_done() + + assert not mock_pub.called