Add include/exclude filter to mqtt_statestream (#10354)

* Add publish filter to mqtt_statestream

* Add tests for include/excludes in mqtt_statestream

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('.', '/') + '/'

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