Hotfix for incompatibility issue between 0.0.3 and 0.0.4

.gitignore

@@ -4,6 +4,7 @@
 __pycache__/
 .idea/
 .coverage
+coverage.xml
 .DS_Store
 *.tmp
 *.log

Makefile

@@ -26,4 +26,7 @@ test:
 	$(POETRY) run pytest $(TEST_FOLDER)
 
 coverage:
-	$(POETRY) run pytest --cov-report xml --cov=custom_components.xiaomi_viomi $(TEST_FOLDER)
+	$(POETRY) run pytest --cov-report xml --cov=custom_components.xiaomi_viomi $(TEST_FOLDER)
+
+coverage-html:
+	$(POETRY) run pytest --cov-report html --cov=custom_components.xiaomi_viomi $(TEST_FOLDER)

custom_components/xiaomi_viomi/config_flow.py

@@ -130,7 +130,8 @@ async def async_step_user(
                 data = existing_entry.data.copy()
                 data[CONF_HOST] = info[CONF_HOST]
                 data[CONF_TOKEN] = info[CONF_TOKEN]
-                data[CONF_NAME] = info[CONF_NAME]
+                data[CONF_NAME] = info[CONF_NAME] or existing_entry.title
+                data[CONF_MODEL] = info[CONF_MODEL]
 
                 self.hass.config_entries.async_update_entry(existing_entry, data=data)
                 await self.hass.config_entries.async_reload(existing_entry.entry_id)

custom_components/xiaomi_viomi/vacuum.py

@@ -3,12 +3,10 @@
 from functools import partial
 from typing import Optional
 
-from homeassistant.components.vacuum import (
-    ATTR_CLEANED_AREA,
-    DOMAIN,
-    STATE_ERROR,
-    StateVacuumEntity,
-)
+from homeassistant.components.vacuum import ATTR_CLEANED_AREA
+from homeassistant.components.vacuum import DOMAIN as PLATFORM_NAME
+from homeassistant.components.vacuum import STATE_ERROR, StateVacuumEntity
+from homeassistant.components.xiaomi_miio import CONF_MODEL
 from homeassistant.components.xiaomi_miio.device import XiaomiMiioEntity
 from homeassistant.config_entries import SOURCE_USER, ConfigEntry
 from homeassistant.const import CONF_HOST, CONF_NAME, CONF_TOKEN, STATE_OFF, STATE_ON
@@ -55,7 +53,7 @@ async def async_setup_platform(
 
     config = await validate_input(hass, raw_config)
     entry = ConfigEntry(
-        domain=DOMAIN,
+        domain=PLATFORM_NAME,
         data=config,
         version=2,
         title=config[CONF_NAME],
@@ -70,9 +68,18 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Xiaomi Viomi config entry."""
-    host = config_entry.data[CONF_HOST]
-    token = config_entry.data[CONF_TOKEN]
-    name = config_entry.data[CONF_NAME]
+
+    if CONF_MODEL not in config_entry.data:
+        data = config_entry.data.copy()
+        data[CONF_NAME] = config_entry.title
+        data[CONF_MODEL] = config_entry.title
+
+        hass.config_entries.async_update_entry(config_entry, data=data)
+
+    host = config_entry.data.get(CONF_HOST)
+    token = config_entry.data.get(CONF_TOKEN)
+    name = config_entry.data.get(CONF_NAME, config_entry.title)
+
     unique_id = config_entry.unique_id
 
     # Create handler

tests/test_setup.py

@@ -1,10 +1,13 @@
 from homeassistant.components.vacuum import DOMAIN, STATE_DOCKED
 from homeassistant.core import HomeAssistant
 from homeassistant.setup import async_setup_component
+from pytest_homeassistant_custom_component.common import MockConfigEntry
 
 from custom_components.xiaomi_viomi import DOMAIN as PLATFORM_NAME
+from custom_components.xiaomi_viomi import async_setup_entry
 from tests import (
     TEST_HOST,
+    TEST_MAC,
     TEST_MODEL,
     TEST_NAME,
     TEST_TOKEN,
@@ -25,13 +28,10 @@ async def test_platform_setup(hass: HomeAssistant):
         ]
     }
 
-    with mocked_viomi_device() as mock_device_send:
+    with mocked_viomi_device():
         await async_setup_component(hass, DOMAIN, config)
-
         await hass.async_block_till_done()
 
-        mock_device_send.reset_mock()
-
         entity_id = get_entity_id()
         state = hass.states.get(entity_id)
 
@@ -51,16 +51,35 @@ async def test_platform_setup_without_name(hass: HomeAssistant):
         ]
     }
 
-    with mocked_viomi_device() as mock_device_send:
+    with mocked_viomi_device():
         await async_setup_component(hass, DOMAIN, config)
-
         await hass.async_block_till_done()
 
-        mock_device_send.reset_mock()
-
         entity_id = get_entity_id(True)
         state = hass.states.get(entity_id)
 
         assert state
         assert state.state == STATE_DOCKED
         assert state.name == TEST_MODEL
+
+
+async def test_setup_with_update_to_004(hass: HomeAssistant):
+    config = MockConfigEntry(
+        domain=PLATFORM_NAME,
+        title=TEST_NAME,
+        data={
+            "host": TEST_HOST,
+            "token": TEST_TOKEN,
+            "mac": TEST_MAC,
+        },
+    )
+
+    with mocked_viomi_device():
+        await async_setup_entry(hass, config)
+        await hass.async_block_till_done()
+
+        entity_id = get_entity_id()
+        state = hass.states.get(entity_id)
+
+        assert state
+        assert state.name == TEST_NAME

tests/test_vacuum.py

@@ -25,13 +25,11 @@
 
 async def test_vacuum_state(hass: HomeAssistant):
     entry = get_mocked_entry()
-    with mocked_viomi_device() as mock_device_send:
+    with mocked_viomi_device():
         entry.add_to_hass(hass)
         await hass.config_entries.async_setup(entry.entry_id)
         await hass.async_block_till_done()
 
-        mock_device_send.reset_mock()
-
         entity_id = get_entity_id()
         state = hass.states.get(entity_id)
 
@@ -178,13 +176,11 @@ async def test_vacuum_fan_speed_service(
 )
 async def test_vacuum_regular_state(hass: HomeAssistant, state_code, state_value):
     entry = get_mocked_entry()
-    with mocked_viomi_device({"run_state": state_code}) as mock_device_send:
+    with mocked_viomi_device({"run_state": state_code}):
         entry.add_to_hass(hass)
         await hass.config_entries.async_setup(entry.entry_id)
         await hass.async_block_till_done()
 
-        mock_device_send.reset_mock()
-
         entity_id = get_entity_id()
         state = hass.states.get(entity_id)
 
@@ -201,13 +197,11 @@ async def test_vacuum_regular_state(hass: HomeAssistant, state_code, state_value
 )
 async def test_vacuum_error_state(hass: HomeAssistant, error_code, error_value):
     entry = get_mocked_entry()
-    with mocked_viomi_device({"err_state": error_code}) as mock_device_send:
+    with mocked_viomi_device({"err_state": error_code}):
         entry.add_to_hass(hass)
         await hass.config_entries.async_setup(entry.entry_id)
         await hass.async_block_till_done()
 
-        mock_device_send.reset_mock()
-
         entity_id = get_entity_id()
         state = hass.states.get(entity_id)