Correct calls to subscription.async_unsubscribe_topics (#19414)

* Correct calls to subscription.async_unsubscribe_topics

* Review comments

* Add testcases

homeassistant/components/alarm_control_panel/mqtt.py

@@ -127,7 +127,8 @@ class MqttAlarm(MqttAvailability, MqttDiscoveryUpdate,
 
     async def async_will_remove_from_hass(self):
         """Unsubscribe when removed."""
-        await subscription.async_unsubscribe_topics(self.hass, self._sub_state)
+        self._sub_state = await subscription.async_unsubscribe_topics(
+            self.hass, self._sub_state)
         await MqttAvailability.async_will_remove_from_hass(self)
 
     @property

homeassistant/components/binary_sensor/mqtt.py

@@ -166,7 +166,8 @@ class MqttBinarySensor(MqttAttributes, MqttAvailability, MqttDiscoveryUpdate,
 
     async def async_will_remove_from_hass(self):
         """Unsubscribe when removed."""
-        await subscription.async_unsubscribe_topics(self.hass, self._sub_state)
+        self._sub_state = await subscription.async_unsubscribe_topics(
+            self.hass, self._sub_state)
         await MqttAttributes.async_will_remove_from_hass(self)
         await MqttAvailability.async_will_remove_from_hass(self)
 

homeassistant/components/climate/mqtt.py

@@ -453,7 +453,8 @@ class MqttClimate(MqttAvailability, MqttDiscoveryUpdate, ClimateDevice):
 
     async def async_will_remove_from_hass(self):
         """Unsubscribe when removed."""
-        await subscription.async_unsubscribe_topics(self.hass, self._sub_state)
+        self._sub_state = await subscription.async_unsubscribe_topics(
+            self.hass, self._sub_state)
         await MqttAvailability.async_will_remove_from_hass(self)
 
     @property

homeassistant/components/cover/mqtt.py

@@ -287,7 +287,8 @@ class MqttCover(MqttAvailability, MqttDiscoveryUpdate, MqttEntityDeviceInfo,
 
     async def async_will_remove_from_hass(self):
         """Unsubscribe when removed."""
-        await subscription.async_unsubscribe_topics(self.hass, self._sub_state)
+        self._sub_state = await subscription.async_unsubscribe_topics(
+            self.hass, self._sub_state)
         await MqttAvailability.async_will_remove_from_hass(self)
 
     @property

homeassistant/components/fan/mqtt.py

@@ -270,7 +270,8 @@ class MqttFan(MqttAvailability, MqttDiscoveryUpdate, MqttEntityDeviceInfo,
 
     async def async_will_remove_from_hass(self):
         """Unsubscribe when removed."""
-        await subscription.async_unsubscribe_topics(self.hass, self._sub_state)
+        self._sub_state = await subscription.async_unsubscribe_topics(
+            self.hass, self._sub_state)
         await MqttAvailability.async_will_remove_from_hass(self)
 
     @property

homeassistant/components/light/mqtt/schema_basic.py

@@ -489,7 +489,8 @@ class MqttLight(MqttAvailability, MqttDiscoveryUpdate, MqttEntityDeviceInfo,
 
     async def async_will_remove_from_hass(self):
         """Unsubscribe when removed."""
-        await subscription.async_unsubscribe_topics(self.hass, self._sub_state)
+        self._sub_state = await subscription.async_unsubscribe_topics(
+            self.hass, self._sub_state)
         await MqttAvailability.async_will_remove_from_hass(self)
 
     @property

homeassistant/components/light/mqtt/schema_json.py

@@ -301,7 +301,8 @@ class MqttLightJson(MqttAvailability, MqttDiscoveryUpdate,
 
     async def async_will_remove_from_hass(self):
         """Unsubscribe when removed."""
-        await subscription.async_unsubscribe_topics(self.hass, self._sub_state)
+        self._sub_state = await subscription.async_unsubscribe_topics(
+            self.hass, self._sub_state)
         await MqttAvailability.async_will_remove_from_hass(self)
 
     @property

homeassistant/components/light/mqtt/schema_template.py

@@ -274,7 +274,8 @@ class MqttTemplate(MqttAvailability, MqttDiscoveryUpdate, MqttEntityDeviceInfo,
 
     async def async_will_remove_from_hass(self):
         """Unsubscribe when removed."""
-        await subscription.async_unsubscribe_topics(self.hass, self._sub_state)
+        self._sub_state = await subscription.async_unsubscribe_topics(
+            self.hass, self._sub_state)
         await MqttAvailability.async_will_remove_from_hass(self)
 
     @property

homeassistant/components/mqtt/__init__.py

@@ -878,7 +878,8 @@ class MqttAttributes(Entity):
     async def async_will_remove_from_hass(self):
         """Unsubscribe when removed."""
         from .subscription import async_unsubscribe_topics
-        await async_unsubscribe_topics(self.hass, self._attributes_sub_state)
+        self._attributes_sub_state = await async_unsubscribe_topics(
+            self.hass, self._attributes_sub_state)
 
     @property
     def device_state_attributes(self):
@@ -947,7 +948,8 @@ class MqttAvailability(Entity):
     async def async_will_remove_from_hass(self):
         """Unsubscribe when removed."""
         from .subscription import async_unsubscribe_topics
-        await async_unsubscribe_topics(self.hass, self._availability_sub_state)
+        self._availability_sub_state = await async_unsubscribe_topics(
+            self.hass, self._availability_sub_state)
 
     @property
     def available(self) -> bool:

homeassistant/components/mqtt/subscription.py

@@ -94,6 +94,4 @@ async def async_subscribe_topics(hass: HomeAssistantType,
 @bind_hass
 async def async_unsubscribe_topics(hass: HomeAssistantType, sub_state: dict):
     """Unsubscribe from all MQTT topics managed by async_subscribe_topics."""
-    await async_subscribe_topics(hass, sub_state, {})
-
-    return sub_state
+    return await async_subscribe_topics(hass, sub_state, {})

homeassistant/components/sensor/mqtt.py

@@ -178,7 +178,8 @@ class MqttSensor(MqttAttributes, MqttAvailability, MqttDiscoveryUpdate,
 
     async def async_will_remove_from_hass(self):
         """Unsubscribe when removed."""
-        await subscription.async_unsubscribe_topics(self.hass, self._sub_state)
+        self._sub_state = await subscription.async_unsubscribe_topics(
+            self.hass, self._sub_state)
         await MqttAttributes.async_will_remove_from_hass(self)
         await MqttAvailability.async_will_remove_from_hass(self)
 

homeassistant/components/switch/mqtt.py

@@ -169,7 +169,8 @@ class MqttSwitch(MqttAvailability, MqttDiscoveryUpdate, MqttEntityDeviceInfo,
 
     async def async_will_remove_from_hass(self):
         """Unsubscribe when removed."""
-        await subscription.async_unsubscribe_topics(self.hass, self._sub_state)
+        self._sub_state = await subscription.async_unsubscribe_topics(
+            self.hass, self._sub_state)
         await MqttAvailability.async_will_remove_from_hass(self)
 
     @property

tests/components/binary_sensor/test_mqtt.py

@@ -1,7 +1,7 @@
 """The tests for the  MQTT binary sensor platform."""
 import json
 import unittest
-from unittest.mock import Mock, patch
+from unittest.mock import ANY, Mock, patch
 from datetime import timedelta
 
 import homeassistant.core as ha
@@ -16,7 +16,7 @@ import homeassistant.util.dt as dt_util
 
 from tests.common import (
     get_test_home_assistant, fire_mqtt_message, async_fire_mqtt_message,
-    fire_time_changed, mock_component, mock_mqtt_component,
+    fire_time_changed, mock_component, mock_mqtt_component, mock_registry,
     async_mock_mqtt_component, MockConfigEntry)
 
 
@@ -457,3 +457,39 @@ async def test_entity_device_info_with_identifier(hass, mqtt_mock):
     assert device.name == 'Beer'
     assert device.model == 'Glass'
     assert device.sw_version == '0.1-beta'
+
+
+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, binary_sensor.DOMAIN, {
+        binary_sensor.DOMAIN: [{
+            'platform': 'mqtt',
+            'name': 'beer',
+            'state_topic': 'test-topic',
+            'availability_topic': 'avty-topic',
+            'unique_id': 'TOTALLY_UNIQUE'
+        }]
+    })
+
+    state = hass.states.get('binary_sensor.beer')
+    assert state is not None
+    assert mock_mqtt.async_subscribe.call_count == 2
+    mock_mqtt.async_subscribe.assert_any_call('test-topic', ANY, 0, 'utf-8')
+    mock_mqtt.async_subscribe.assert_any_call('avty-topic', ANY, 0, 'utf-8')
+    mock_mqtt.async_subscribe.reset_mock()
+
+    registry.async_update_entity(
+        'binary_sensor.beer', new_entity_id='binary_sensor.milk')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    state = hass.states.get('binary_sensor.beer')
+    assert state is None
+
+    state = hass.states.get('binary_sensor.milk')
+    assert state is not None
+    assert mock_mqtt.async_subscribe.call_count == 2
+    mock_mqtt.async_subscribe.assert_any_call('test-topic', ANY, 0, 'utf-8')
+    mock_mqtt.async_subscribe.assert_any_call('avty-topic', ANY, 0, 'utf-8')

tests/components/cover/test_mqtt.py

@@ -1,6 +1,7 @@
 """The tests for the MQTT cover platform."""
 import json
 import unittest
+from unittest.mock import ANY
 
 from homeassistant.components import cover, mqtt
 from homeassistant.components.cover import (ATTR_POSITION, ATTR_TILT_POSITION)
@@ -16,7 +17,8 @@ from homeassistant.setup import setup_component, async_setup_component
 
 from tests.common import (
     get_test_home_assistant, mock_mqtt_component, async_fire_mqtt_message,
-    fire_mqtt_message, MockConfigEntry, async_mock_mqtt_component)
+    fire_mqtt_message, MockConfigEntry, async_mock_mqtt_component,
+    mock_registry)
 
 
 class TestCoverMQTT(unittest.TestCase):
@@ -1156,3 +1158,38 @@ async def test_entity_device_info_with_identifier(hass, mqtt_mock):
     assert device.name == 'Beer'
     assert device.model == 'Glass'
     assert device.sw_version == '0.1-beta'
+
+
+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, cover.DOMAIN, {
+        cover.DOMAIN: [{
+            'platform': 'mqtt',
+            'name': 'beer',
+            'state_topic': 'test-topic',
+            'availability_topic': 'avty-topic',
+            'unique_id': 'TOTALLY_UNIQUE'
+        }]
+    })
+
+    state = hass.states.get('cover.beer')
+    assert state is not None
+    assert mock_mqtt.async_subscribe.call_count == 2
+    mock_mqtt.async_subscribe.assert_any_call('test-topic', ANY, 0, 'utf-8')
+    mock_mqtt.async_subscribe.assert_any_call('avty-topic', ANY, 0, 'utf-8')
+    mock_mqtt.async_subscribe.reset_mock()
+
+    registry.async_update_entity('cover.beer', new_entity_id='cover.milk')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    state = hass.states.get('cover.beer')
+    assert state is None
+
+    state = hass.states.get('cover.milk')
+    assert state is not None
+    assert mock_mqtt.async_subscribe.call_count == 2
+    mock_mqtt.async_subscribe.assert_any_call('test-topic', ANY, 0, 'utf-8')
+    mock_mqtt.async_subscribe.assert_any_call('avty-topic', ANY, 0, 'utf-8')

tests/components/fan/test_mqtt.py

@@ -1,5 +1,6 @@
 """Test MQTT fans."""
 import json
+from unittest.mock import ANY
 
 from homeassistant.setup import async_setup_component
 from homeassistant.components import fan
@@ -7,7 +8,7 @@ from homeassistant.components.mqtt.discovery import async_start
 from homeassistant.const import ATTR_ASSUMED_STATE, STATE_UNAVAILABLE
 
 from tests.common import async_fire_mqtt_message, MockConfigEntry, \
-    async_mock_mqtt_component
+    async_mock_mqtt_component, mock_registry
 
 
 async def test_fail_setup_if_no_command_topic(hass, mqtt_mock):
@@ -224,3 +225,39 @@ async def test_entity_device_info_with_identifier(hass, mqtt_mock):
     assert device.name == 'Beer'
     assert device.model == 'Glass'
     assert device.sw_version == '0.1-beta'
+
+
+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, fan.DOMAIN, {
+        fan.DOMAIN: [{
+            'platform': 'mqtt',
+            'name': 'beer',
+            'state_topic': 'test-topic',
+            'command_topic': 'command-topic',
+            'availability_topic': 'avty-topic',
+            'unique_id': 'TOTALLY_UNIQUE'
+        }]
+    })
+
+    state = hass.states.get('fan.beer')
+    assert state is not None
+    assert mock_mqtt.async_subscribe.call_count == 2
+    mock_mqtt.async_subscribe.assert_any_call('test-topic', ANY, 0, 'utf-8')
+    mock_mqtt.async_subscribe.assert_any_call('avty-topic', ANY, 0, 'utf-8')
+    mock_mqtt.async_subscribe.reset_mock()
+
+    registry.async_update_entity('fan.beer', new_entity_id='fan.milk')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    state = hass.states.get('fan.beer')
+    assert state is None
+
+    state = hass.states.get('fan.milk')
+    assert state is not None
+    assert mock_mqtt.async_subscribe.call_count == 2
+    mock_mqtt.async_subscribe.assert_any_call('test-topic', ANY, 0, 'utf-8')
+    mock_mqtt.async_subscribe.assert_any_call('avty-topic', ANY, 0, 'utf-8')

tests/components/light/test_mqtt.py

@@ -155,7 +155,7 @@ light:
 """
 import json
 from unittest import mock
-from unittest.mock import patch
+from unittest.mock import ANY, patch
 
 from homeassistant.setup import async_setup_component
 from homeassistant.const import (
@@ -165,8 +165,8 @@ from homeassistant.components.mqtt.discovery import async_start
 import homeassistant.core as ha
 
 from tests.common import (
-    assert_setup_component, async_fire_mqtt_message,
-    async_mock_mqtt_component, mock_coro, MockConfigEntry)
+    assert_setup_component, async_fire_mqtt_message, async_mock_mqtt_component,
+    mock_coro, MockConfigEntry, mock_registry)
 from tests.components.light import common
 
 
@@ -1180,3 +1180,39 @@ async def test_entity_device_info_with_identifier(hass, mqtt_mock):
     assert device.name == 'Beer'
     assert device.model == 'Glass'
     assert device.sw_version == '0.1-beta'
+
+
+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, light.DOMAIN, {
+        light.DOMAIN: [{
+            'platform': 'mqtt',
+            'name': 'beer',
+            'state_topic': 'test-topic',
+            'command_topic': 'command-topic',
+            'availability_topic': 'avty-topic',
+            'unique_id': 'TOTALLY_UNIQUE'
+        }]
+    })
+
+    state = hass.states.get('light.beer')
+    assert state is not None
+    assert mock_mqtt.async_subscribe.call_count == 2
+    mock_mqtt.async_subscribe.assert_any_call('test-topic', ANY, 0, 'utf-8')
+    mock_mqtt.async_subscribe.assert_any_call('avty-topic', ANY, 0, 'utf-8')
+    mock_mqtt.async_subscribe.reset_mock()
+
+    registry.async_update_entity('light.beer', new_entity_id='light.milk')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    state = hass.states.get('light.beer')
+    assert state is None
+
+    state = hass.states.get('light.milk')
+    assert state is not None
+    assert mock_mqtt.async_subscribe.call_count == 2
+    mock_mqtt.async_subscribe.assert_any_call('test-topic', ANY, 0, 'utf-8')
+    mock_mqtt.async_subscribe.assert_any_call('avty-topic', ANY, 0, 'utf-8')

tests/components/light/test_mqtt_json.py

@@ -88,7 +88,7 @@ light:
   brightness_scale: 99
 """
 import json
-from unittest.mock import patch
+from unittest.mock import ANY, patch
 
 from homeassistant.setup import async_setup_component
 from homeassistant.const import (
@@ -100,7 +100,7 @@ import homeassistant.core as ha
 
 from tests.common import (
     mock_coro, async_fire_mqtt_message, async_mock_mqtt_component,
-    MockConfigEntry)
+    MockConfigEntry, mock_registry)
 
 
 async def test_fail_setup_if_no_command_topic(hass, mqtt_mock):
@@ -677,3 +677,40 @@ async def test_entity_device_info_with_identifier(hass, mqtt_mock):
     assert device.name == 'Beer'
     assert device.model == 'Glass'
     assert device.sw_version == '0.1-beta'
+
+
+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, light.DOMAIN, {
+        light.DOMAIN: [{
+            'platform': 'mqtt',
+            'name': 'beer',
+            'schema': 'json',
+            'state_topic': 'test-topic',
+            'command_topic': 'command-topic',
+            'availability_topic': 'avty-topic',
+            'unique_id': 'TOTALLY_UNIQUE'
+        }]
+    })
+
+    state = hass.states.get('light.beer')
+    assert state is not None
+    assert mock_mqtt.async_subscribe.call_count == 2
+    mock_mqtt.async_subscribe.assert_any_call('test-topic', ANY, 0, 'utf-8')
+    mock_mqtt.async_subscribe.assert_any_call('avty-topic', ANY, 0, 'utf-8')
+    mock_mqtt.async_subscribe.reset_mock()
+
+    registry.async_update_entity('light.beer', new_entity_id='light.milk')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    state = hass.states.get('light.beer')
+    assert state is None
+
+    state = hass.states.get('light.milk')
+    assert state is not None
+    assert mock_mqtt.async_subscribe.call_count == 2
+    mock_mqtt.async_subscribe.assert_any_call('test-topic', ANY, 0, 'utf-8')
+    mock_mqtt.async_subscribe.assert_any_call('avty-topic', ANY, 0, 'utf-8')

tests/components/light/test_mqtt_template.py

@@ -27,7 +27,7 @@ If your light doesn't support white value feature, omit `white_value_template`.
 If your light doesn't support RGB feature, omit `(red|green|blue)_template`.
 """
 import json
-from unittest.mock import patch
+from unittest.mock import ANY, patch
 
 from homeassistant.setup import async_setup_component
 from homeassistant.const import (
@@ -38,7 +38,7 @@ import homeassistant.core as ha
 
 from tests.common import (
     async_fire_mqtt_message, assert_setup_component, mock_coro,
-    async_mock_mqtt_component, MockConfigEntry)
+    async_mock_mqtt_component, MockConfigEntry, mock_registry)
 
 
 async def test_setup_fails(hass, mqtt_mock):
@@ -632,3 +632,42 @@ async def test_entity_device_info_with_identifier(hass, mqtt_mock):
     assert device.name == 'Beer'
     assert device.model == 'Glass'
     assert device.sw_version == '0.1-beta'
+
+
+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, light.DOMAIN, {
+        light.DOMAIN: [{
+            'platform': 'mqtt',
+            'name': 'beer',
+            'schema': 'template',
+            'state_topic': 'test-topic',
+            'command_topic': 'command-topic',
+            'command_on_template': 'on,{{ transition }}',
+            'command_off_template': 'off,{{ transition|d }}',
+            'availability_topic': 'avty-topic',
+            'unique_id': 'TOTALLY_UNIQUE'
+        }]
+    })
+
+    state = hass.states.get('light.beer')
+    assert state is not None
+    assert mock_mqtt.async_subscribe.call_count == 2
+    mock_mqtt.async_subscribe.assert_any_call('test-topic', ANY, 0, 'utf-8')
+    mock_mqtt.async_subscribe.assert_any_call('avty-topic', ANY, 0, 'utf-8')
+    mock_mqtt.async_subscribe.reset_mock()
+
+    registry.async_update_entity('light.beer', new_entity_id='light.milk')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    state = hass.states.get('light.beer')
+    assert state is None
+
+    state = hass.states.get('light.milk')
+    assert state is not None
+    assert mock_mqtt.async_subscribe.call_count == 2
+    mock_mqtt.async_subscribe.assert_any_call('test-topic', ANY, 0, 'utf-8')
+    mock_mqtt.async_subscribe.assert_any_call('avty-topic', ANY, 0, 'utf-8')

tests/components/sensor/test_mqtt.py

@@ -3,7 +3,7 @@ import json
 import unittest
 
 from datetime import timedelta, datetime
-from unittest.mock import patch
+from unittest.mock import ANY, patch
 
 import homeassistant.core as ha
 from homeassistant.setup import setup_component, async_setup_component
@@ -15,7 +15,7 @@ import homeassistant.util.dt as dt_util
 
 from tests.common import mock_mqtt_component, fire_mqtt_message, \
     assert_setup_component, async_fire_mqtt_message, \
-    async_mock_mqtt_component, MockConfigEntry
+    async_mock_mqtt_component, MockConfigEntry, mock_registry
 from tests.common import get_test_home_assistant, mock_component
 
 
@@ -542,3 +542,38 @@ async def test_entity_device_info_with_identifier(hass, mqtt_mock):
     assert device.name == 'Beer'
     assert device.model == 'Glass'
     assert device.sw_version == '0.1-beta'
+
+
+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, sensor.DOMAIN, {
+        sensor.DOMAIN: [{
+            'platform': 'mqtt',
+            'name': 'beer',
+            'state_topic': 'test-topic',
+            'availability_topic': 'avty-topic',
+            'unique_id': 'TOTALLY_UNIQUE'
+        }]
+    })
+
+    state = hass.states.get('sensor.beer')
+    assert state is not None
+    assert mock_mqtt.async_subscribe.call_count == 2
+    mock_mqtt.async_subscribe.assert_any_call('test-topic', ANY, 0, 'utf-8')
+    mock_mqtt.async_subscribe.assert_any_call('avty-topic', ANY, 0, 'utf-8')
+    mock_mqtt.async_subscribe.reset_mock()
+
+    registry.async_update_entity('sensor.beer', new_entity_id='sensor.milk')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    state = hass.states.get('sensor.beer')
+    assert state is None
+
+    state = hass.states.get('sensor.milk')
+    assert state is not None
+    assert mock_mqtt.async_subscribe.call_count == 2
+    mock_mqtt.async_subscribe.assert_any_call('test-topic', ANY, 0, 'utf-8')
+    mock_mqtt.async_subscribe.assert_any_call('avty-topic', ANY, 0, 'utf-8')

tests/components/switch/test_mqtt.py

@@ -2,6 +2,7 @@
 import json
 from asynctest import patch
 import pytest
+from unittest.mock import ANY
 
 from homeassistant.setup import async_setup_component
 from homeassistant.const import STATE_ON, STATE_OFF, STATE_UNAVAILABLE,\
@@ -12,7 +13,7 @@ from homeassistant.components.mqtt.discovery import async_start
 
 from tests.common import (
     mock_coro, async_mock_mqtt_component, async_fire_mqtt_message,
-    MockConfigEntry)
+    MockConfigEntry, mock_registry)
 from tests.components.switch import common
 
 
@@ -434,3 +435,39 @@ async def test_entity_device_info_with_identifier(hass, mqtt_mock):
     assert device.name == 'Beer'
     assert device.model == 'Glass'
     assert device.sw_version == '0.1-beta'
+
+
+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, switch.DOMAIN, {
+        switch.DOMAIN: [{
+            'platform': 'mqtt',
+            'name': 'beer',
+            'state_topic': 'test-topic',
+            'command_topic': 'command-topic',
+            'availability_topic': 'avty-topic',
+            'unique_id': 'TOTALLY_UNIQUE'
+        }]
+    })
+
+    state = hass.states.get('switch.beer')
+    assert state is not None
+    assert mock_mqtt.async_subscribe.call_count == 2
+    mock_mqtt.async_subscribe.assert_any_call('test-topic', ANY, 0, 'utf-8')
+    mock_mqtt.async_subscribe.assert_any_call('avty-topic', ANY, 0, 'utf-8')
+    mock_mqtt.async_subscribe.reset_mock()
+
+    registry.async_update_entity('switch.beer', new_entity_id='switch.milk')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    state = hass.states.get('switch.beer')
+    assert state is None
+
+    state = hass.states.get('switch.milk')
+    assert state is not None
+    assert mock_mqtt.async_subscribe.call_count == 2
+    mock_mqtt.async_subscribe.assert_any_call('test-topic', ANY, 0, 'utf-8')
+    mock_mqtt.async_subscribe.assert_any_call('avty-topic', ANY, 0, 'utf-8')