Add include/exclude filter to mqtt_statestream (#10354)

* Add publish filter to mqtt_statestream

* Add tests for include/excludes in mqtt_statestream
pull/10465/head
Matt White 2017-11-08 08:54:12 -06:00 committed by Pascal Vizeli
parent 2f0920e4fb
commit 148a7ddda9
2 changed files with 254 additions and 4 deletions

View File

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

View File

@ -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