Add include/exclude filter to mqtt_statestream (#10354)
* Add publish filter to mqtt_statestream * Add tests for include/excludes in mqtt_statestreampull/10465/head
parent
2f0920e4fb
commit
148a7ddda9
|
@ -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('.', '/') + '/'
|
||||
|
|
|
@ -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
|
||||
|
|
Loading…
Reference in New Issue