Add discovery update support to MQTT camera (#20529)

homeassistant/components/mqtt/camera.py

@@ -12,8 +12,10 @@ import voluptuous as vol
 
 from homeassistant.components import camera, mqtt
 from homeassistant.components.camera import PLATFORM_SCHEMA, Camera
-from homeassistant.components.mqtt import CONF_UNIQUE_ID
+from homeassistant.components.mqtt import (
-from homeassistant.components.mqtt.discovery import MQTT_DISCOVERY_NEW
+    ATTR_DISCOVERY_HASH, CONF_UNIQUE_ID, MqttDiscoveryUpdate, subscription)
+from homeassistant.components.mqtt.discovery import (
+    MQTT_DISCOVERY_NEW, clear_discovery_hash)
 from homeassistant.const import CONF_NAME
 from homeassistant.core import callback
 from homeassistant.helpers import config_validation as cv
@@ -37,43 +39,79 @@ PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({
 async def async_setup_platform(hass: HomeAssistantType, config: ConfigType,
                                async_add_entities, discovery_info=None):
     """Set up MQTT camera through configuration.yaml."""
-    await _async_setup_entity(hass, config, async_add_entities)
+    await _async_setup_entity(config, async_add_entities)
 
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
     """Set up MQTT camera dynamically through MQTT discovery."""
     async def async_discover(discovery_payload):
         """Discover and add a MQTT camera."""
+        try:
+            discovery_hash = discovery_payload[ATTR_DISCOVERY_HASH]
             config = PLATFORM_SCHEMA(discovery_payload)
-        await _async_setup_entity(hass, config, async_add_entities)
+            await _async_setup_entity(config, async_add_entities,
+                                      discovery_hash)
+        except Exception:
+            if discovery_hash:
+                clear_discovery_hash(hass, discovery_hash)
+            raise
 
     async_dispatcher_connect(
         hass, MQTT_DISCOVERY_NEW.format(camera.DOMAIN, 'mqtt'),
         async_discover)
 
 
-async def _async_setup_entity(hass, config, async_add_entities):
+async def _async_setup_entity(config, async_add_entities, discovery_hash=None):
     """Set up the MQTT Camera."""
-    async_add_entities([MqttCamera(
+    async_add_entities([MqttCamera(config, discovery_hash)])
-        config.get(CONF_NAME),
-        config.get(CONF_UNIQUE_ID),
-        config.get(CONF_TOPIC)
-    )])
 
 
-class MqttCamera(Camera):
+class MqttCamera(MqttDiscoveryUpdate, Camera):
     """representation of a MQTT camera."""
 
-    def __init__(self, name, unique_id, topic):
+    def __init__(self, config, discovery_hash):
         """Initialize the MQTT Camera."""
-        super().__init__()
+        self._config = config
+        self._unique_id = config.get(CONF_UNIQUE_ID)
+        self._sub_state = None
 
-        self._name = name
-        self._unique_id = unique_id
-        self._topic = topic
         self._qos = 0
         self._last_image = None
 
+        Camera.__init__(self)
+        MqttDiscoveryUpdate.__init__(self, discovery_hash,
+                                     self.discovery_update)
+
+    async def async_added_to_hass(self):
+        """Subscribe MQTT events."""
+        await super().async_added_to_hass()
+        await self._subscribe_topics()
+
+    async def discovery_update(self, discovery_payload):
+        """Handle updated discovery message."""
+        config = PLATFORM_SCHEMA(discovery_payload)
+        self._config = config
+        await self._subscribe_topics()
+        self.async_schedule_update_ha_state()
+
+    async def _subscribe_topics(self):
+        """(Re)Subscribe to topics."""
+        @callback
+        def message_received(topic, payload, qos):
+            """Handle new MQTT messages."""
+            self._last_image = payload
+
+        self._sub_state = await subscription.async_subscribe_topics(
+            self.hass, self._sub_state,
+            {'state_topic': {'topic': self._config.get(CONF_TOPIC),
+                             'msg_callback': message_received,
+                             'qos': self._qos}})
+
+    async def async_will_remove_from_hass(self):
+        """Unsubscribe when removed."""
+        self._sub_state = await subscription.async_unsubscribe_topics(
+            self.hass, self._sub_state)
+
     @asyncio.coroutine
     def async_camera_image(self):
         """Return image response."""
@@ -82,19 +120,9 @@ class MqttCamera(Camera):
     @property
     def name(self):
         """Return the name of this camera."""
-        return self._name
+        return self._config.get(CONF_NAME)
 
     @property
     def unique_id(self):
         """Return a unique ID."""
         return self._unique_id
-
-    async def async_added_to_hass(self):
-        """Subscribe MQTT events."""
-        @callback
-        def message_received(topic, payload, qos):
-            """Handle new MQTT messages."""
-            self._last_image = payload
-
-        await mqtt.async_subscribe(
-            self.hass, self._topic, message_received, self._qos, None)

tests/components/mqtt/test_camera.py

@@ -1,9 +1,14 @@
 """The tests for mqtt camera component."""
 import asyncio
+from unittest.mock import ANY
 
+from homeassistant.components import camera, mqtt
+from homeassistant.components.mqtt.discovery import async_start
 from homeassistant.setup import async_setup_component
 
-from tests.common import async_fire_mqtt_message, async_mock_mqtt_component
+from tests.common import (
+    MockConfigEntry, async_fire_mqtt_message, async_mock_mqtt_component,
+    mock_registry)
 
 
 @asyncio.coroutine
@@ -51,3 +56,129 @@ def test_unique_id(hass):
     async_fire_mqtt_message(hass, 'test-topic', 'payload')
     yield from hass.async_block_till_done()
     assert len(hass.states.async_all()) == 1
+
+
+async def test_discovery_removal_camera(hass, mqtt_mock, caplog):
+    """Test removal of discovered camera."""
+    entry = MockConfigEntry(domain=mqtt.DOMAIN)
+    await async_start(hass, 'homeassistant', {}, entry)
+
+    data = (
+        '{ "name": "Beer",'
+        '  "topic": "test_topic"}'
+    )
+
+    async_fire_mqtt_message(hass, 'homeassistant/camera/bla/config',
+                            data)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    state = hass.states.get('camera.beer')
+    assert state is not None
+    assert state.name == 'Beer'
+
+    async_fire_mqtt_message(hass, 'homeassistant/camera/bla/config',
+                            '')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    state = hass.states.get('camera.beer')
+    assert state is None
+
+
+async def test_discovery_update_camera(hass, mqtt_mock, caplog):
+    """Test update of discovered camera."""
+    entry = MockConfigEntry(domain=mqtt.DOMAIN)
+    await async_start(hass, 'homeassistant', {}, entry)
+
+    data1 = (
+        '{ "name": "Beer",'
+        '  "topic": "test_topic"}'
+    )
+    data2 = (
+        '{ "name": "Milk",'
+        '  "topic": "test_topic"}'
+    )
+
+    async_fire_mqtt_message(hass, 'homeassistant/camera/bla/config',
+                            data1)
+    await hass.async_block_till_done()
+
+    state = hass.states.get('camera.beer')
+    assert state is not None
+    assert state.name == 'Beer'
+
+    async_fire_mqtt_message(hass, 'homeassistant/camera/bla/config',
+                            data2)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    state = hass.states.get('camera.beer')
+    assert state is not None
+    assert state.name == 'Milk'
+    state = hass.states.get('camera.milk')
+    assert state is None
+
+
+async def test_discovery_broken(hass, mqtt_mock, caplog):
+    """Test handling of bad discovery message."""
+    entry = MockConfigEntry(domain=mqtt.DOMAIN)
+    await async_start(hass, 'homeassistant', {}, entry)
+
+    data1 = (
+        '{ "name": "Beer" }'
+    )
+    data2 = (
+        '{ "name": "Milk",'
+        '  "topic": "test_topic"}'
+    )
+
+    async_fire_mqtt_message(hass, 'homeassistant/camera/bla/config',
+                            data1)
+    await hass.async_block_till_done()
+
+    state = hass.states.get('camera.beer')
+    assert state is None
+
+    async_fire_mqtt_message(hass, 'homeassistant/camera/bla/config',
+                            data2)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    state = hass.states.get('camera.milk')
+    assert state is not None
+    assert state.name == 'Milk'
+    state = hass.states.get('camera.beer')
+    assert state is None
+
+
+async def test_entity_id_update(hass, mqtt_mock):
+    """Test MQTT subscriptions are managed when entity_id is updated."""
+    registry = mock_registry(hass, {})
+    mock_mqtt = await async_mock_mqtt_component(hass)
+    assert await async_setup_component(hass, camera.DOMAIN, {
+        camera.DOMAIN: [{
+            'platform': 'mqtt',
+            'name': 'beer',
+            'topic': 'test-topic',
+            'unique_id': 'TOTALLY_UNIQUE'
+        }]
+    })
+
+    state = hass.states.get('camera.beer')
+    assert state is not None
+    assert mock_mqtt.async_subscribe.call_count == 1
+    mock_mqtt.async_subscribe.assert_any_call('test-topic', ANY, 0, 'utf-8')
+    mock_mqtt.async_subscribe.reset_mock()
+
+    registry.async_update_entity('camera.beer', new_entity_id='camera.milk')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    state = hass.states.get('camera.beer')
+    assert state is None
+
+    state = hass.states.get('camera.milk')
+    assert state is not None
+    assert mock_mqtt.async_subscribe.call_count == 1
+    mock_mqtt.async_subscribe.assert_any_call('test-topic', ANY, 0, 'utf-8')