diff --git a/homeassistant/components/mqtt_statestream.py b/homeassistant/components/mqtt_statestream.py index 2b68394b160..8469cb3b334 100644 --- a/homeassistant/components/mqtt_statestream.py +++ b/homeassistant/components/mqtt_statestream.py @@ -12,14 +12,19 @@ from homeassistant.const import MATCH_ALL from homeassistant.core import callback from homeassistant.components.mqtt import valid_publish_topic from homeassistant.helpers.event import async_track_state_change +import homeassistant.helpers.config_validation as cv CONF_BASE_TOPIC = 'base_topic' +CONF_PUBLISH_ATTRIBUTES = 'publish_attributes' +CONF_PUBLISH_TIMESTAMPS = 'publish_timestamps' DEPENDENCIES = ['mqtt'] DOMAIN = 'mqtt_statestream' CONFIG_SCHEMA = vol.Schema({ DOMAIN: vol.Schema({ - vol.Required(CONF_BASE_TOPIC): valid_publish_topic + 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 }) }, extra=vol.ALLOW_EXTRA) @@ -29,6 +34,8 @@ def async_setup(hass, config): """Set up the MQTT state feed.""" conf = config.get(DOMAIN, {}) base_topic = conf.get(CONF_BASE_TOPIC) + publish_attributes = conf.get(CONF_PUBLISH_ATTRIBUTES) + publish_timestamps = conf.get(CONF_PUBLISH_TIMESTAMPS) if not base_topic.endswith('/'): base_topic = base_topic + '/' @@ -38,8 +45,28 @@ def async_setup(hass, config): return payload = new_state.state - topic = base_topic + entity_id.replace('.', '/') + '/state' - hass.components.mqtt.async_publish(topic, payload, 1, True) + mybase = base_topic + entity_id.replace('.', '/') + '/' + hass.components.mqtt.async_publish(mybase + 'state', payload, 1, True) + + if publish_timestamps: + if new_state.last_updated: + hass.components.mqtt.async_publish( + mybase + 'last_updated', + new_state.last_updated.isoformat(), + 1, + True) + if new_state.last_changed: + hass.components.mqtt.async_publish( + mybase + 'last_changed', + new_state.last_changed.isoformat(), + 1, + True) + + if publish_attributes: + for key, val in new_state.attributes.items(): + if val: + hass.components.mqtt.async_publish(mybase + key, + val, 1, True) async_track_state_change(hass, MATCH_ALL, _state_publisher) return True diff --git a/tests/components/test_mqtt_statestream.py b/tests/components/test_mqtt_statestream.py index cbd7838effe..802d62bfdd1 100644 --- a/tests/components/test_mqtt_statestream.py +++ b/tests/components/test_mqtt_statestream.py @@ -1,5 +1,5 @@ """The tests for the MQTT statestream component.""" -from unittest.mock import patch +from unittest.mock import ANY, call, patch from homeassistant.setup import setup_component import homeassistant.components.mqtt_statestream as statestream @@ -24,11 +24,17 @@ class TestMqttStateStream(object): """Stop everything that was started.""" self.hass.stop() - def add_statestream(self, base_topic=None): + def add_statestream(self, base_topic=None, publish_attributes=None, + publish_timestamps=None): """Add a mqtt_statestream component.""" config = {} if base_topic: config['base_topic'] = base_topic + if publish_attributes: + config['publish_attributes'] = publish_attributes + if publish_timestamps: + config['publish_timestamps'] = publish_timestamps + print("Publishing timestamps") return setup_component(self.hass, statestream.DOMAIN, { statestream.DOMAIN: config}) @@ -36,10 +42,14 @@ class TestMqttStateStream(object): """Setup should fail if no base_topic is set.""" assert self.add_statestream() is False - def test_setup_succeeds(self): + def test_setup_succeeds_without_attributes(self): """"Test the success of the setup with a valid base_topic.""" assert self.add_statestream(base_topic='pub') + def test_setup_succeeds_with_attributes(self): + """"Test setup with a valid base_topic and publish_attributes.""" + assert self.add_statestream(base_topic='pub', publish_attributes=True) + @patch('homeassistant.components.mqtt.async_publish') @patch('homeassistant.core.dt_util.utcnow') def test_state_changed_event_sends_message(self, mock_utcnow, mock_pub): @@ -60,6 +70,77 @@ class TestMqttStateStream(object): 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) + mock_pub.assert_called_with(self.hass, 'pub/fake/entity/state', 'on', + 1, True) + assert mock_pub.called + + @patch('homeassistant.components.mqtt.async_publish') + @patch('homeassistant.core.dt_util.utcnow') + def test_state_changed_event_sends_message_and_timestamp( + self, + mock_utcnow, + mock_pub): + """"Test the sending of a message and timestamps if event changed.""" + e_id = 'another.entity' + base_topic = 'pub' + + # Add the statestream component for publishing state updates + assert self.add_statestream(base_topic=base_topic, + publish_attributes=None, + publish_timestamps=True) + 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(e_id, 'on')) + self.hass.block_till_done() + + # Make sure 'on' was published to pub/fake/entity/state + calls = [ + call.async_publish(self.hass, 'pub/another/entity/state', 'on', 1, + True), + call.async_publish(self.hass, 'pub/another/entity/last_changed', + ANY, 1, True), + call.async_publish(self.hass, 'pub/another/entity/last_updated', + ANY, 1, True), + ] + + 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_attr_sends_message(self, mock_utcnow, mock_pub): + """"Test the sending of a new message if attribute changed.""" + e_id = 'fake.entity' + base_topic = 'pub' + + # Add the statestream component for publishing state updates + assert self.add_statestream(base_topic=base_topic, + publish_attributes=True) + 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() + + test_attributes = {"testing": "YES"} + + # Set a state of an entity + mock_state_change_event(self.hass, State(e_id, 'off', + attributes=test_attributes)) + self.hass.block_till_done() + + # Make sure 'on' was published to pub/fake/entity/state + calls = [ + call.async_publish(self.hass, 'pub/fake/entity/state', 'off', 1, + True), + call.async_publish(self.hass, 'pub/fake/entity/testing', 'YES', + 1, True) + ] + + mock_pub.assert_has_calls(calls, any_order=True) assert mock_pub.called