From e341694bceb338b29942bd4123a1e264fe417cbd Mon Sep 17 00:00:00 2001
From: Artem Poliukhovych <apoliu@softserveinc.com>
Date: Tue, 25 Jan 2022 23:54:21 +0200
Subject: [PATCH 1/8] Add platform configuration

---
 README.md                                     | 13 +++++++++++-
 custom_components/xiaomi_viomi/__init__.py    | 12 +++++++++++
 custom_components/xiaomi_viomi/config_flow.py |  8 +++----
 custom_components/xiaomi_viomi/const.py       |  3 ---
 custom_components/xiaomi_viomi/vacuum.py      | 21 +++++++++++++++++--
 tests/test_config_flow.py                     |  2 +-
 6 files changed, 48 insertions(+), 11 deletions(-)

diff --git a/README.md b/README.md
index 8eefbc3..6a5b843 100644
--- a/README.md
+++ b/README.md
@@ -14,7 +14,18 @@ Xiaomi Viomi vacuum robot integration for Home Assistant.
 ### HACS
 Install it through HACS by adding this as a custom repository: https://github.com/nergal/homeassistant-vacuum-viomi, go to the integrations page in your configuration, click on the `Add Integration` button in the bottom right corner of a screen, and search for `Xiaomi Viomi Vacuum`.
 
-Manual installation currently not supported.
+### Manual
+Copy contents of `custom_components` folder to your Home Assistant `config/custom_components` folder. Restart Home Assistant, and then the integration can be added and configured through the native integration setup. If you don't see it in the native integrations list, press Ctrl+F5 to refresh the browser while you're on that page and retry.
+
+Also you may add the manual configuration to `configuration.yaml` file, like the example below:
+
+```
+vacuum:
+  - platform: xiaomi_viomi
+    host: 192.168.1.1
+    token: !secret vacuum
+    name: Vacuum V8
+```
 
 ## Tested models
 | Model | Device ID | Aliases | Status |
diff --git a/custom_components/xiaomi_viomi/__init__.py b/custom_components/xiaomi_viomi/__init__.py
index 90e3e89..d456fb6 100644
--- a/custom_components/xiaomi_viomi/__init__.py
+++ b/custom_components/xiaomi_viomi/__init__.py
@@ -1,11 +1,23 @@
 """Xiaomi Viomi integration."""
 import asyncio
 
+import voluptuous as vol
+from homeassistant.components.vacuum import PLATFORM_SCHEMA
 from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import CONF_HOST, CONF_NAME, CONF_TOKEN, DEVICE_DEFAULT_NAME
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers import config_validation as cv
 
 PLATFORMS = ["vacuum"]
 
+PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
+    {
+        vol.Optional(CONF_NAME, default=DEVICE_DEFAULT_NAME): cv.string,
+        vol.Required(CONF_TOKEN): cv.string,
+        vol.Required(CONF_HOST): cv.string,
+    }
+)
+
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Xiaomi Viomi from a config entry."""
diff --git a/custom_components/xiaomi_viomi/config_flow.py b/custom_components/xiaomi_viomi/config_flow.py
index 358207a..6c6925b 100644
--- a/custom_components/xiaomi_viomi/config_flow.py
+++ b/custom_components/xiaomi_viomi/config_flow.py
@@ -5,7 +5,7 @@
 import voluptuous as vol
 from construct.core import ChecksumError
 from homeassistant import config_entries
-from homeassistant.const import CONF_HOST, CONF_TOKEN
+from homeassistant.const import CONF_HOST, CONF_MAC, CONF_NAME, CONF_TOKEN
 from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import FlowResult
 from homeassistant.exceptions import ConfigEntryAuthFailed, HomeAssistantError
@@ -14,7 +14,7 @@
 from miio.device import DeviceInfo
 from miio.integrations.vacuum.viomi.viomivacuum import ViomiVacuum
 
-from .const import CONF_MAC, CONF_MODEL, DOMAIN
+from .const import DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -79,7 +79,7 @@ async def validate_input(hass: HomeAssistant, data: Dict[str, Any]) -> Dict[str,
     return {
         CONF_HOST: data[CONF_HOST],
         CONF_TOKEN: data[CONF_TOKEN],
-        CONF_MODEL: hub.device_info.model,
+        CONF_NAME: hub.device_info.model,
         CONF_MAC: format_mac(hub.device_info.mac_address),
     }
 
@@ -121,7 +121,7 @@ async def async_step_user(
                 await self.hass.config_entries.async_reload(existing_entry.entry_id)
                 return self.async_abort(reason="already_configured")
 
-            return self.async_create_entry(title=info[CONF_MODEL], data=info)
+            return self.async_create_entry(title=info[CONF_NAME], data=info)
 
         return self.async_show_form(
             step_id="user", data_schema=DEVICE_CONFIG, errors=errors
diff --git a/custom_components/xiaomi_viomi/const.py b/custom_components/xiaomi_viomi/const.py
index ac49d5b..ce31263 100644
--- a/custom_components/xiaomi_viomi/const.py
+++ b/custom_components/xiaomi_viomi/const.py
@@ -1,9 +1,6 @@
 """Constants for Xiaomi Viomi integration."""
 
 DOMAIN = "xiaomi_viomi"
-CONF_DEVICE = "device"
-CONF_MODEL = "model"
-CONF_MAC = "mac"
 CONF_FLOW_TYPE = "config_flow_device"
 
 MODELS_VACUUM = ["viomi.vacuum.v8"]
diff --git a/custom_components/xiaomi_viomi/vacuum.py b/custom_components/xiaomi_viomi/vacuum.py
index 9a70d69..6e01741 100644
--- a/custom_components/xiaomi_viomi/vacuum.py
+++ b/custom_components/xiaomi_viomi/vacuum.py
@@ -22,6 +22,8 @@
 from homeassistant.components.xiaomi_miio.device import XiaomiMiioEntity
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, CONF_TOKEN, STATE_OFF, STATE_ON
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from miio import DeviceException
 from miio.integrations.vacuum.viomi.viomivacuum import (
     ViomiVacuum,
@@ -67,15 +69,21 @@
 }
 
 ERRORS_FALSE_POSITIVE = (
-    0,  # Sleeping and not charging
+    0,  # Sleeping and not charging,
+    514,  # ? Stuck
     2103,  # Charging
     2104,  # ? Returning
     2105,  # Fully charged
+    2108,  # ? Stuck
     2110,  # ? Cleaning
 )
 
 
-async def async_setup_entry(hass, config_entry: ConfigEntry, async_add_entities):
+async def async_setup_platform(
+    _hass: HomeAssistant,
+    config_entry: ConfigEntry,
+    async_add_entities: AddEntitiesCallback,
+) -> None:
     """Set up the Xiaomi Viomi vacuum cleaner robot from a config entry."""
     entities = []
 
@@ -94,6 +102,15 @@ async def async_setup_entry(hass, config_entry: ConfigEntry, async_add_entities)
     async_add_entities(entities, update_before_add=True)
 
 
+async def async_setup_entry(
+    hass: HomeAssistant,
+    config_entry: ConfigEntry,
+    async_add_entities: AddEntitiesCallback,
+) -> None:
+    """Set up the Xiaomi Viomi config entry."""
+    await async_setup_platform(hass, config_entry, async_add_entities)
+
+
 class ViomiVacuumIntegration(XiaomiMiioEntity, StateVacuumEntity):
     """Xiaomi Viomi integration handler."""
 
diff --git a/tests/test_config_flow.py b/tests/test_config_flow.py
index 732d5f9..ecf6037 100644
--- a/tests/test_config_flow.py
+++ b/tests/test_config_flow.py
@@ -49,7 +49,7 @@ async def test_form(hass: HomeAssistant) -> None:
 
     result_data = {
         "host": "1.1.1.1",
-        "model": "Name of the device",
+        "name": "Name of the device",
         "token": "ffffffffffffffffffffffffffffffff",
         "mac": "f2:ff:ff:ff:ff:ff",
     }

From 942229e5da72d213dfcbc23c4e84cbf346601b2e Mon Sep 17 00:00:00 2001
From: Artem Poliukhovych <nergal@users.noreply.github.com>
Date: Fri, 4 Feb 2022 02:31:43 +0200
Subject: [PATCH 2/8] Add unit tests, platform setup and locate

---
 .gitignore                                    |   2 +
 Makefile                                      |   2 +-
 custom_components/xiaomi_viomi/__init__.py    |   5 +
 custom_components/xiaomi_viomi/config_flow.py |   8 +-
 custom_components/xiaomi_viomi/const.py       |  60 +++++-
 custom_components/xiaomi_viomi/vacuum.py      | 176 +++++++++---------
 tests/__init__.py                             |  89 +++++++++
 tests/test_config_flow.py                     |  35 ++--
 tests/test_setup.py                           |  40 ++++
 tests/test_vacuum.py                          | 121 ++++++++++++
 10 files changed, 435 insertions(+), 103 deletions(-)
 create mode 100644 tests/test_setup.py
 create mode 100644 tests/test_vacuum.py

diff --git a/.gitignore b/.gitignore
index 5e9e046..d36ce3d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,5 +3,7 @@
 .vscode/
 __pycache__/
 .idea/
+.coverage
+.DS_Store
 *.tmp
 *.log
\ No newline at end of file
diff --git a/Makefile b/Makefile
index cc3d0cd..1eceb6c 100644
--- a/Makefile
+++ b/Makefile
@@ -26,4 +26,4 @@ test:
 	$(POETRY) run pytest $(TEST_FOLDER)
 
 coverage:
-	$(POETRY) run pytest --cov-report term $(TEST_FOLDER)
\ No newline at end of file
+	$(POETRY) run pytest --cov-report term --cov=custom_components.xiaomi_viomi $(TEST_FOLDER)
\ No newline at end of file
diff --git a/custom_components/xiaomi_viomi/__init__.py b/custom_components/xiaomi_viomi/__init__.py
index d456fb6..2488f44 100644
--- a/custom_components/xiaomi_viomi/__init__.py
+++ b/custom_components/xiaomi_viomi/__init__.py
@@ -3,16 +3,21 @@
 
 import voluptuous as vol
 from homeassistant.components.vacuum import PLATFORM_SCHEMA
+from homeassistant.components.xiaomi_miio import CONF_MODEL
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, CONF_NAME, CONF_TOKEN, DEVICE_DEFAULT_NAME
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import config_validation as cv
+from miio.integrations.vacuum.viomi.viomivacuum import SUPPORTED_MODELS
+
+from .const import DOMAIN  # noqa: F401
 
 PLATFORMS = ["vacuum"]
 
 PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
     {
         vol.Optional(CONF_NAME, default=DEVICE_DEFAULT_NAME): cv.string,
+        vol.Required(CONF_MODEL): cv.ensure_list(SUPPORTED_MODELS),
         vol.Required(CONF_TOKEN): cv.string,
         vol.Required(CONF_HOST): cv.string,
     }
diff --git a/custom_components/xiaomi_viomi/config_flow.py b/custom_components/xiaomi_viomi/config_flow.py
index 6c6925b..959ecf3 100644
--- a/custom_components/xiaomi_viomi/config_flow.py
+++ b/custom_components/xiaomi_viomi/config_flow.py
@@ -5,6 +5,7 @@
 import voluptuous as vol
 from construct.core import ChecksumError
 from homeassistant import config_entries
+from homeassistant.components.xiaomi_miio import CONF_MODEL
 from homeassistant.const import CONF_HOST, CONF_MAC, CONF_NAME, CONF_TOKEN
 from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import FlowResult
@@ -12,7 +13,7 @@
 from homeassistant.helpers.device_registry import format_mac
 from miio import DeviceException
 from miio.device import DeviceInfo
-from miio.integrations.vacuum.viomi.viomivacuum import ViomiVacuum
+from miio.integrations.vacuum.viomi.viomivacuum import SUPPORTED_MODELS, ViomiVacuum
 
 from .const import DOMAIN
 
@@ -21,7 +22,9 @@
 DEVICE_CONFIG = vol.Schema(
     {
         vol.Required(CONF_HOST): str,
+        vol.Required(CONF_NAME): str,
         vol.Required(CONF_TOKEN): vol.All(str, vol.Length(min=32, max=32)),
+        vol.Required(CONF_MODEL): vol.In(SUPPORTED_MODELS),
     }
 )
 
@@ -79,7 +82,8 @@ async def validate_input(hass: HomeAssistant, data: Dict[str, Any]) -> Dict[str,
     return {
         CONF_HOST: data[CONF_HOST],
         CONF_TOKEN: data[CONF_TOKEN],
-        CONF_NAME: hub.device_info.model,
+        CONF_MODEL: hub.device_info.model,
+        CONF_NAME: data[CONF_NAME],
         CONF_MAC: format_mac(hub.device_info.mac_address),
     }
 
diff --git a/custom_components/xiaomi_viomi/const.py b/custom_components/xiaomi_viomi/const.py
index ce31263..b80a28d 100644
--- a/custom_components/xiaomi_viomi/const.py
+++ b/custom_components/xiaomi_viomi/const.py
@@ -1,10 +1,24 @@
 """Constants for Xiaomi Viomi integration."""
 
+from homeassistant.components.vacuum import (
+    STATE_CLEANING,
+    STATE_DOCKED,
+    STATE_IDLE,
+    STATE_RETURNING,
+    SUPPORT_BATTERY,
+    SUPPORT_FAN_SPEED,
+    SUPPORT_LOCATE,
+    SUPPORT_PAUSE,
+    SUPPORT_RETURN_HOME,
+    SUPPORT_SEND_COMMAND,
+    SUPPORT_START,
+    SUPPORT_STATE,
+    SUPPORT_STOP,
+)
+
 DOMAIN = "xiaomi_viomi"
 CONF_FLOW_TYPE = "config_flow_device"
 
-MODELS_VACUUM = ["viomi.vacuum.v8"]
-
 DEVICE_PROPERTIES = [
     "battary_life",
     "box_type",
@@ -43,3 +57,45 @@
     "hypa_hours",
     "hypa_life",
 ]
+
+ATTR_CLEANING_TIME = "cleaning_time"
+ATTR_DO_NOT_DISTURB = "do_not_disturb"
+ATTR_DO_NOT_DISTURB_START = "do_not_disturb_start"
+ATTR_DO_NOT_DISTURB_END = "do_not_disturb_end"
+ATTR_MAIN_BRUSH_LEFT = "main_brush_left"
+ATTR_SIDE_BRUSH_LEFT = "side_brush_left"
+ATTR_FILTER_LEFT = "filter_left"
+ATTR_MOP_LEFT = "mop_left"
+ATTR_ERROR = "error"
+ATTR_STATUS = "status"
+ATTR_MOP_ATTACHED = "mop_attached"
+
+ERRORS_FALSE_POSITIVE = (
+    0,  # Sleeping and not charging,
+    2103,  # Charging
+    2104,  # ? Returning
+    2105,  # Fully charged
+    2110,  # ? Cleaning
+)
+
+SUPPORT_VIOMI = (
+    SUPPORT_PAUSE
+    | SUPPORT_STOP
+    | SUPPORT_RETURN_HOME
+    | SUPPORT_FAN_SPEED
+    | SUPPORT_BATTERY
+    | SUPPORT_SEND_COMMAND
+    | SUPPORT_LOCATE
+    | SUPPORT_STATE
+    | SUPPORT_START
+)
+
+STATE_CODE_TO_STATE = {
+    0: STATE_IDLE,  # IdleNotDocked
+    1: STATE_IDLE,  # Idle
+    2: STATE_IDLE,  # Idle2
+    3: STATE_CLEANING,  # Cleaning
+    4: STATE_RETURNING,  # Returning
+    5: STATE_DOCKED,  # Docked
+    6: STATE_CLEANING,  # VacuumingAndMopping
+}
diff --git a/custom_components/xiaomi_viomi/vacuum.py b/custom_components/xiaomi_viomi/vacuum.py
index 6e01741..782d5e0 100644
--- a/custom_components/xiaomi_viomi/vacuum.py
+++ b/custom_components/xiaomi_viomi/vacuum.py
@@ -1,124 +1,105 @@
 """Xiaomi Viomi integration."""
 import logging
 from functools import partial
+from typing import Optional
 
 from homeassistant.components.vacuum import (
     ATTR_CLEANED_AREA,
-    STATE_CLEANING,
-    STATE_DOCKED,
+    DOMAIN,
     STATE_ERROR,
-    STATE_IDLE,
-    STATE_RETURNING,
-    SUPPORT_BATTERY,
-    SUPPORT_FAN_SPEED,
-    SUPPORT_PAUSE,
-    SUPPORT_RETURN_HOME,
-    SUPPORT_SEND_COMMAND,
-    SUPPORT_START,
-    SUPPORT_STATE,
-    SUPPORT_STOP,
     StateVacuumEntity,
 )
 from homeassistant.components.xiaomi_miio.device import XiaomiMiioEntity
-from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_HOST, CONF_TOKEN, STATE_OFF, STATE_ON
+from homeassistant.config_entries import SOURCE_USER, ConfigEntry
+from homeassistant.const import CONF_HOST, CONF_NAME, CONF_TOKEN, STATE_OFF, STATE_ON
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
 from miio import DeviceException
+from miio.click_common import command
 from miio.integrations.vacuum.viomi.viomivacuum import (
     ViomiVacuum,
     ViomiVacuumSpeed,
     ViomiVacuumStatus,
 )
 
-from .const import DEVICE_PROPERTIES
-
-_LOGGER = logging.getLogger(__name__)
-
-ATTR_CLEANING_TIME = "cleaning_time"
-ATTR_DO_NOT_DISTURB = "do_not_disturb"
-ATTR_DO_NOT_DISTURB_START = "do_not_disturb_start"
-ATTR_DO_NOT_DISTURB_END = "do_not_disturb_end"
-ATTR_MAIN_BRUSH_LEFT = "main_brush_left"
-ATTR_SIDE_BRUSH_LEFT = "side_brush_left"
-ATTR_FILTER_LEFT = "filter_left"
-ATTR_MOP_LEFT = "mop_left"
-ATTR_ERROR = "error"
-ATTR_STATUS = "status"
-ATTR_MOP_ATTACHED = "mop_attached"
-
-SUPPORT_VIOMI = (
-    SUPPORT_STATE
-    | SUPPORT_PAUSE
-    | SUPPORT_STOP
-    | SUPPORT_RETURN_HOME
-    | SUPPORT_FAN_SPEED
-    | SUPPORT_SEND_COMMAND
-    | SUPPORT_BATTERY
-    | SUPPORT_START
+from .config_flow import validate_input
+from .const import (
+    ATTR_CLEANING_TIME,
+    ATTR_DO_NOT_DISTURB,
+    ATTR_DO_NOT_DISTURB_END,
+    ATTR_DO_NOT_DISTURB_START,
+    ATTR_ERROR,
+    ATTR_FILTER_LEFT,
+    ATTR_MAIN_BRUSH_LEFT,
+    ATTR_MOP_ATTACHED,
+    ATTR_MOP_LEFT,
+    ATTR_SIDE_BRUSH_LEFT,
+    ATTR_STATUS,
+    DEVICE_PROPERTIES,
+    ERRORS_FALSE_POSITIVE,
+    STATE_CODE_TO_STATE,
+    SUPPORT_VIOMI,
 )
 
-STATE_CODE_TO_STATE = {
-    0: STATE_IDLE,  # IdleNotDocked
-    1: STATE_IDLE,  # Idle
-    2: STATE_IDLE,  # Idle2
-    3: STATE_CLEANING,  # Cleaning
-    4: STATE_RETURNING,  # Returning
-    5: STATE_DOCKED,  # Docked
-    6: STATE_CLEANING,  # VacuumingAndMopping
-}
-
-ERRORS_FALSE_POSITIVE = (
-    0,  # Sleeping and not charging,
-    514,  # ? Stuck
-    2103,  # Charging
-    2104,  # ? Returning
-    2105,  # Fully charged
-    2108,  # ? Stuck
-    2110,  # ? Cleaning
-)
+_LOGGER = logging.getLogger(__name__)
 
 
 async def async_setup_platform(
-    _hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    hass: HomeAssistant,
+    raw_config: ConfigType,
     async_add_entities: AddEntitiesCallback,
+    discovery_info: Optional[DiscoveryInfoType] = None,
 ) -> None:
     """Set up the Xiaomi Viomi vacuum cleaner robot from a config entry."""
-    entities = []
 
+    config = await validate_input(hass, raw_config)
+    entry = ConfigEntry(
+        domain=DOMAIN,
+        data=config,
+        version=2,
+        title=raw_config[CONF_NAME],
+        source=SOURCE_USER,
+    )
+    await async_setup_entry(hass, entry, async_add_entities)
+
+
+async def async_setup_entry(
+    hass: HomeAssistant,
+    config_entry: ConfigEntry,
+    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.title
+    name = config_entry.data[CONF_NAME]
     unique_id = config_entry.unique_id
 
     # Create handler
     _LOGGER.debug("Initializing viomi with host %s (token %s...)", host, token[:5])
-    vacuum = ViomiVacuum(host, token)
+    vacuum = PatchedViomiVacuum(ip=host, token=token)
 
     viomi = ViomiVacuumIntegration(name, vacuum, config_entry, unique_id)
-    entities.append(viomi)
-
-    async_add_entities(entities, update_before_add=True)
+    async_add_entities([viomi], update_before_add=True)
 
 
-async def async_setup_entry(
-    hass: HomeAssistant,
-    config_entry: ConfigEntry,
-    async_add_entities: AddEntitiesCallback,
-) -> None:
-    """Set up the Xiaomi Viomi config entry."""
-    await async_setup_platform(hass, config_entry, async_add_entities)
+class PatchedViomiVacuum(ViomiVacuum):
+    @command()
+    def locate(self):
+        """Locate a device."""
+        self.send("set_resetpos", [1])
 
 
 class ViomiVacuumIntegration(XiaomiMiioEntity, StateVacuumEntity):
     """Xiaomi Viomi integration handler."""
 
+    _device: PatchedViomiVacuum
+
     def __init__(self, name, device, entry, unique_id):
         """Initialize the Xiaomi vacuum cleaner robot handler."""
         super().__init__(name, device, entry, unique_id)
 
-        self.vacuum_state = None
+        self.vacuum_state: Optional[ViomiVacuumStatus] = None
         self._available = False
 
         self.consumable_state = None
@@ -127,7 +108,7 @@ def __init__(self, name, device, entry, unique_id):
         self._fan_speeds_reverse = None
 
     @property
-    def state(self):
+    def state(self) -> Optional[str]:
         """Return the status of the vacuum cleaner."""
         if self.vacuum_state is not None:
             # The vacuum reverts back to an idle state after erroring out.
@@ -148,7 +129,8 @@ def state(self):
                     self.vacuum_state.state,
                     self.vacuum_state.state.value,
                 )
-                return None
+
+        return None
 
     @property
     def battery_level(self):
@@ -221,18 +203,19 @@ def available(self) -> bool:
         return self._available
 
     @property
-    def supported_features(self):
+    def supported_features(self) -> int:
         """Flag vacuum cleaner robot features that are supported."""
         return SUPPORT_VIOMI
 
-    def _got_error(self):
-        error_code = self.vacuum_state.error_code
-        return error_code and error_code not in ERRORS_FALSE_POSITIVE
+    def _got_error(self) -> bool:
+        error_code = self.vacuum_state.error_code if self.vacuum_state else None
+        return bool(error_code and error_code not in ERRORS_FALSE_POSITIVE)
 
-    def _get_status(self):
-        return STATE_CODE_TO_STATE[int(self.vacuum_state.state.value)]
+    def _get_status(self) -> str:
+        state_value = self.vacuum_state.state.value if self.vacuum_state else None
+        return STATE_CODE_TO_STATE[int(state_value)]
 
-    def _get_device_status(self):
+    def _get_device_status(self) -> ViomiVacuumStatus:
         """Override of miio's device.status() because of bug."""
         result = {}
         for prop in DEVICE_PROPERTIES:
@@ -268,6 +251,18 @@ async def _try_command(self, mask_error, func, *args, **kwargs):
             _LOGGER.error(mask_error, exc)
             return False
 
+    async def async_turn_on(self, **kwargs):
+        """Start or resume the cleaning task."""
+        await self.async_start()
+
+    async def async_turn_off(self, **kwargs):
+        """Stop the cleaning task."""
+        await self.async_stop()
+
+    async def async_toggle(self, **kwargs):
+        """Start or pause depending on current state."""
+        await self.async_start_pause()
+
     async def async_start(self):
         """Start or resume the cleaning task."""
         await self._try_command("Unable to start the vacuum: %s", self._device.start)
@@ -276,10 +271,25 @@ async def async_pause(self):
         """Pause the cleaning task."""
         await self._try_command("Unable to set start/pause: %s", self._device.pause)
 
+    async def async_start_pause(self):
+        """Start or pause depending on current state."""
+        if self.vacuum_state.is_on:
+            await self.async_pause()
+        else:
+            await self.async_start()
+
     async def async_stop(self, **kwargs):
         """Stop the vacuum cleaner."""
         await self._try_command("Unable to stop: %s", self._device.stop)
 
+    async def async_locate(self, **kwargs):
+        """Locate the vacuum cleaner."""
+        await self._try_command("Unable to locate: %s", self._device.locate)
+
+    async def async_clean_spot(self, **kwargs):
+        """Suppress NotImplementedError for clean spot capability."""
+        pass
+
     async def async_set_fan_speed(self, fan_speed, **kwargs):
         """Set fan speed."""
         if fan_speed in self._fan_speeds:
diff --git a/tests/__init__.py b/tests/__init__.py
index b694d5a..f0491ec 100644
--- a/tests/__init__.py
+++ b/tests/__init__.py
@@ -1 +1,90 @@
 """Tests for the Xiaomi Viomi integration."""
+from typing import Any
+from unittest.mock import patch
+
+from homeassistant.components.vacuum import DOMAIN
+from pytest_homeassistant_custom_component.common import MockConfigEntry
+
+from custom_components.xiaomi_viomi.const import DOMAIN as CUSTOM_DOMAIN
+
+TEST_HOST = "1.1.1.1"
+TEST_MAC = "f2:ff:ff:ff:ff:ff"
+TEST_MODEL = "viomi.vacuum.v8"
+TEST_TOKEN = "ffffffffffffffffffffffffffffffff"
+TEST_NAME = "mocked_vacuum"
+
+MOCKING_SEND_METHOD = "miio.integrations.vacuum.viomi.viomivacuum.ViomiVacuum.send"
+
+MOCKED_DEVICE_STATE = {
+    "battary_life": 100,
+    "box_type": 1,
+    "err_state": 2105,
+    "has_map": 1,
+    "has_newmap": 1,
+    "hw_info": "1.0.3",
+    "is_charge": 0,
+    "is_mop": 0,
+    "is_work": 1,
+    "light_state": 1,
+    "mode": 0,
+    "mop_type": 0,
+    "order_time": "0",
+    "remember_map": 1,
+    "repeat_state": 0,
+    "run_state": 5,
+    "s_area": 11.96,
+    "s_time": 20,
+    "start_time": 0,
+    "suction_grade": 0,
+    "sw_info": "3.5.3_0017",
+    "v_state": 10,
+    "water_grade": 12,
+    "zone_data": "0",
+}
+
+MOCKED_DEVICE_INFO = {
+    "model": TEST_MODEL,
+    "mac": TEST_MAC,
+}
+
+
+def get_entity_id() -> str:
+    return DOMAIN + "." + TEST_NAME
+
+
+def get_mocked_entry():
+    return MockConfigEntry(
+        domain=CUSTOM_DOMAIN,
+        title=TEST_NAME,
+        data={
+            "host": TEST_HOST,
+            "token": TEST_TOKEN,
+            "name": TEST_NAME,
+            "mac": TEST_MAC,
+            "model": TEST_MODEL,
+        },
+    )
+
+
+def mocked_viomi_device():
+    def _device_mock_method(command: str, parameters: Any = None):
+        # Request for getting device state
+        if command == "get_prop" and parameters:
+            property_name = parameters[0]
+            if property_name in MOCKED_DEVICE_STATE:
+                return [MOCKED_DEVICE_STATE[property_name]]
+
+            return [None]
+        # Request for getting state of consumables
+        elif command == "get_consumables":
+            return [0] * 5
+        # Request DND configuration
+        elif command == "get_notdisturb":
+            return [0] * 5
+        # Request information about the hardware
+        elif command == "miIO.info":
+            return MOCKED_DEVICE_INFO
+
+        return None
+
+    return patch(MOCKING_SEND_METHOD, wraps=_device_mock_method)
diff --git a/tests/test_config_flow.py b/tests/test_config_flow.py
index ecf6037..bda52e0 100644
--- a/tests/test_config_flow.py
+++ b/tests/test_config_flow.py
@@ -8,6 +8,7 @@
 
 from custom_components.xiaomi_viomi.config_flow import CannotConnect, InvalidAuth
 from custom_components.xiaomi_viomi.const import DOMAIN
+from tests import TEST_HOST, TEST_MAC, TEST_MODEL, TEST_NAME, TEST_TOKEN
 
 
 async def test_form(hass: HomeAssistant) -> None:
@@ -20,10 +21,7 @@ async def test_form(hass: HomeAssistant) -> None:
     assert result["type"] == RESULT_TYPE_FORM
     assert result["errors"] is None
 
-    device_info_mock = {
-        "model": "Name of the device",
-        "mac_address": "F2:FF:FF:FF:FF:FF",
-    }
+    device_info_mock = {"mac_address": TEST_MAC, "model": TEST_MODEL}
 
     with patch(
         "custom_components.xiaomi_viomi.config_flow.ViomiDeviceHub.async_device_is_connectable",
@@ -41,21 +39,24 @@ async def test_form(hass: HomeAssistant) -> None:
         result2 = await hass.config_entries.flow.async_configure(
             result["flow_id"],
             {
-                "host": "1.1.1.1",
-                "token": "ffffffffffffffffffffffffffffffff",
+                "host": TEST_HOST,
+                "token": TEST_TOKEN,
+                "model": TEST_MODEL,
+                "name": TEST_NAME,
             },
         )
         await hass.async_block_till_done()
 
     result_data = {
-        "host": "1.1.1.1",
-        "name": "Name of the device",
-        "token": "ffffffffffffffffffffffffffffffff",
-        "mac": "f2:ff:ff:ff:ff:ff",
+        "host": TEST_HOST,
+        "name": TEST_NAME,
+        "token": TEST_TOKEN,
+        "model": TEST_MODEL,
+        "mac": TEST_MAC,
     }
 
     assert result2["type"] == RESULT_TYPE_CREATE_ENTRY
-    assert result2["title"] == "Name of the device"
+    assert result2["title"] == TEST_NAME
     assert result2["data"] == result_data
     assert len(mock_setup_entry.mock_calls) == 1
 
@@ -73,8 +74,10 @@ async def test_form_invalid_auth(hass: HomeAssistant) -> None:
         result2 = await hass.config_entries.flow.async_configure(
             result["flow_id"],
             {
-                "host": "1.1.1.1",
-                "token": "ffffffffffffffffffffffffffffffff",
+                "host": TEST_HOST,
+                "token": TEST_TOKEN,
+                "model": TEST_MODEL,
+                "name": TEST_NAME,
             },
         )
 
@@ -95,8 +98,10 @@ async def test_form_cannot_connect(hass: HomeAssistant) -> None:
         result2 = await hass.config_entries.flow.async_configure(
             result["flow_id"],
             {
-                "host": "1.1.1.1",
-                "token": "ffffffffffffffffffffffffffffffff",
+                "host": TEST_HOST,
+                "token": TEST_TOKEN,
+                "model": TEST_MODEL,
+                "name": TEST_NAME,
             },
         )
 
diff --git a/tests/test_setup.py b/tests/test_setup.py
new file mode 100644
index 0000000..c6b8929
--- /dev/null
+++ b/tests/test_setup.py
@@ -0,0 +1,40 @@
+from homeassistant.components.vacuum import DOMAIN, STATE_DOCKED
+from homeassistant.core import HomeAssistant
+from homeassistant.setup import async_setup_component
+
+from custom_components.xiaomi_viomi import DOMAIN as PLATFORM_NAME
+from tests import (
+    TEST_HOST,
+    TEST_MODEL,
+    TEST_NAME,
+    TEST_TOKEN,
+    get_entity_id,
+    mocked_viomi_device,
+)
+
+
+async def test_platform_setup(hass: HomeAssistant):
+    config = {
+        DOMAIN: [
+            {
+                "platform": PLATFORM_NAME,
+                "host": TEST_HOST,
+                "token": TEST_TOKEN,
+                "name": TEST_NAME,
+                "model": TEST_MODEL,
+            }
+        ]
+    }
+
+    with mocked_viomi_device() as mock_device_send:
+        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)
+
+        assert state
+        assert state.state == STATE_DOCKED
diff --git a/tests/test_vacuum.py b/tests/test_vacuum.py
new file mode 100644
index 0000000..101f88c
--- /dev/null
+++ b/tests/test_vacuum.py
@@ -0,0 +1,121 @@
+"""Test sensor for simple integration."""
+import pytest
+from homeassistant.components.vacuum import (
+    DOMAIN,
+    SERVICE_LOCATE,
+    SERVICE_PAUSE,
+    SERVICE_RETURN_TO_BASE,
+    SERVICE_SEND_COMMAND,
+    SERVICE_SET_FAN_SPEED,
+    SERVICE_START,
+    SERVICE_START_PAUSE,
+    SERVICE_STOP,
+    STATE_DOCKED,
+)
+from homeassistant.const import SERVICE_TOGGLE, SERVICE_TURN_OFF, SERVICE_TURN_ON
+from homeassistant.core import HomeAssistant
+from miio.integrations.vacuum.viomi.viomivacuum import ViomiVacuumSpeed
+
+from custom_components.xiaomi_viomi.const import SUPPORT_VIOMI as SUPPORT_FEATURES
+from tests import get_entity_id, get_mocked_entry, mocked_viomi_device
+
+
+async def test_vacuum_state(hass: HomeAssistant):
+    entry = get_mocked_entry()
+    with mocked_viomi_device() as mock_device_send:
+        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)
+
+        assert state
+        assert state.state == STATE_DOCKED
+        assert state.attributes["battery_level"] == 100
+        assert state.attributes["fan_speed"] == ViomiVacuumSpeed.Silent.name
+        assert state.attributes["supported_features"] == SUPPORT_FEATURES
+
+
+test_service_data = [
+    (SERVICE_TURN_ON, "set_mode_withroom", [0, 1, 0]),
+    (SERVICE_TURN_OFF, "set_mode", [0, 0]),
+    (SERVICE_TOGGLE, "set_mode_withroom", [0, 1, 0]),
+    (SERVICE_START_PAUSE, "set_mode_withroom", [0, 1, 0]),
+    (SERVICE_START, "set_mode_withroom", [0, 1, 0]),
+    (SERVICE_PAUSE, "set_mode", [0, 2]),
+    (SERVICE_RETURN_TO_BASE, "set_charge", [1]),
+    (SERVICE_LOCATE, "set_resetpos", [1]),
+    (SERVICE_STOP, "set_mode", [0, 0]),
+]
+
+
+@pytest.mark.parametrize("service,method,parameters", test_service_data)
+async def test_vacuum_service(hass: HomeAssistant, service, method, parameters):
+    entry = get_mocked_entry()
+    with mocked_viomi_device() as mock_device_send:
+        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()
+
+        await hass.services.async_call(
+            DOMAIN, service, {"entity_id": entity_id}, blocking=True
+        )
+        mock_device_send.assert_any_call(method, parameters)
+
+
+async def test_vacuum_send_command_service(hass: HomeAssistant):
+    entry = get_mocked_entry()
+    with mocked_viomi_device() as mock_device_send:
+        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()
+
+        await hass.services.async_call(
+            DOMAIN,
+            SERVICE_SEND_COMMAND,
+            {"command": "test_command", "entity_id": entity_id},
+            blocking=True,
+        )
+
+        mock_device_send.assert_any_call("test_command", None)
+
+
+@pytest.mark.parametrize(
+    "speed",
+    [
+        ViomiVacuumSpeed.Silent,
+        ViomiVacuumSpeed.Medium,
+        ViomiVacuumSpeed.Standard,
+        ViomiVacuumSpeed.Turbo,
+    ],
+)
+async def test_vacuum_fan_speed_service(hass: HomeAssistant, speed: ViomiVacuumSpeed):
+    entry = get_mocked_entry()
+    with mocked_viomi_device() as mock_device_send:
+        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()
+
+        await hass.services.async_call(
+            DOMAIN,
+            SERVICE_SET_FAN_SPEED,
+            {"entity_id": entity_id, "fan_speed": speed.value},
+            blocking=True,
+        )
+
+        mock_device_send.assert_any_call("set_suction", [speed.value])

From 8965cc35db963f16d87a40d6afd50ae9c6c8a974 Mon Sep 17 00:00:00 2001
From: Artem Poliukhovych <nergal@users.noreply.github.com>
Date: Fri, 4 Feb 2022 13:57:29 +0200
Subject: [PATCH 3/8] Update unit tests and add coverage metrics

---
 .github/workflows/workflow.yaml               |   6 +-
 codecov.yml                                   |   0
 custom_components/xiaomi_viomi/__init__.py    |   3 -
 custom_components/xiaomi_viomi/config_flow.py |  21 +++-
 custom_components/xiaomi_viomi/strings.json   |   6 +-
 .../xiaomi_viomi/translations/en.json         |   4 +-
 .../xiaomi_viomi/translations/ru.json         |   4 +-
 .../xiaomi_viomi/translations/uk.json         |   4 +-
 custom_components/xiaomi_viomi/vacuum.py      |   8 +-
 tests/__init__.py                             |  16 +--
 tests/test_config_flow.py                     | 118 +++++++++++++-----
 tests/test_setup.py                           |  28 ++++-
 tests/test_vacuum.py                          | 112 +++++++++++++++--
 13 files changed, 260 insertions(+), 70 deletions(-)
 create mode 100644 codecov.yml

diff --git a/.github/workflows/workflow.yaml b/.github/workflows/workflow.yaml
index b504174..c65b352 100644
--- a/.github/workflows/workflow.yaml
+++ b/.github/workflows/workflow.yaml
@@ -13,7 +13,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: [3.8, 3.9]
+        python-version: [3.8, 3.9, 3.10]
         os: [ubuntu-latest, macos-latest]
     runs-on: ${{ matrix.os }}
     steps:
@@ -30,7 +30,9 @@ jobs:
       - name: Run lint checks
         run: make lint
       - name: Run unit tests
-        run: make test
+        run: make coverage
+      - name: Upload Coverage to Codecov
+        uses: codecov/codecov-action@v1
 
   analyze:
     name: Analyze CodeQL
diff --git a/codecov.yml b/codecov.yml
new file mode 100644
index 0000000..e69de29
diff --git a/custom_components/xiaomi_viomi/__init__.py b/custom_components/xiaomi_viomi/__init__.py
index 2488f44..dcbdf56 100644
--- a/custom_components/xiaomi_viomi/__init__.py
+++ b/custom_components/xiaomi_viomi/__init__.py
@@ -3,12 +3,10 @@
 
 import voluptuous as vol
 from homeassistant.components.vacuum import PLATFORM_SCHEMA
-from homeassistant.components.xiaomi_miio import CONF_MODEL
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, CONF_NAME, CONF_TOKEN, DEVICE_DEFAULT_NAME
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import config_validation as cv
-from miio.integrations.vacuum.viomi.viomivacuum import SUPPORTED_MODELS
 
 from .const import DOMAIN  # noqa: F401
 
@@ -17,7 +15,6 @@
 PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
     {
         vol.Optional(CONF_NAME, default=DEVICE_DEFAULT_NAME): cv.string,
-        vol.Required(CONF_MODEL): cv.ensure_list(SUPPORTED_MODELS),
         vol.Required(CONF_TOKEN): cv.string,
         vol.Required(CONF_HOST): cv.string,
     }
diff --git a/custom_components/xiaomi_viomi/config_flow.py b/custom_components/xiaomi_viomi/config_flow.py
index 959ecf3..e9c5473 100644
--- a/custom_components/xiaomi_viomi/config_flow.py
+++ b/custom_components/xiaomi_viomi/config_flow.py
@@ -6,14 +6,21 @@
 from construct.core import ChecksumError
 from homeassistant import config_entries
 from homeassistant.components.xiaomi_miio import CONF_MODEL
-from homeassistant.const import CONF_HOST, CONF_MAC, CONF_NAME, CONF_TOKEN
+from homeassistant.const import (
+    CONF_HOST,
+    CONF_MAC,
+    CONF_NAME,
+    CONF_PLATFORM,
+    CONF_TOKEN,
+    DEVICE_DEFAULT_NAME,
+)
 from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import FlowResult
 from homeassistant.exceptions import ConfigEntryAuthFailed, HomeAssistantError
 from homeassistant.helpers.device_registry import format_mac
 from miio import DeviceException
 from miio.device import DeviceInfo
-from miio.integrations.vacuum.viomi.viomivacuum import SUPPORTED_MODELS, ViomiVacuum
+from miio.integrations.vacuum.viomi.viomivacuum import ViomiVacuum
 
 from .const import DOMAIN
 
@@ -22,9 +29,8 @@
 DEVICE_CONFIG = vol.Schema(
     {
         vol.Required(CONF_HOST): str,
-        vol.Required(CONF_NAME): str,
         vol.Required(CONF_TOKEN): vol.All(str, vol.Length(min=32, max=32)),
-        vol.Required(CONF_MODEL): vol.In(SUPPORTED_MODELS),
+        vol.Optional(CONF_NAME, default=DEVICE_DEFAULT_NAME): str,
     }
 )
 
@@ -79,11 +85,15 @@ async def validate_input(hass: HomeAssistant, data: Dict[str, Any]) -> Dict[str,
     if not await hub.async_device_is_connectable(data[CONF_HOST], data[CONF_TOKEN]):
         raise InvalidAuth
 
+    DEVICE_CONFIG.extend({vol.Optional(CONF_PLATFORM): str})(data)
+
+    name = data[CONF_NAME] if CONF_NAME in data else hub.device_info.model
+
     return {
         CONF_HOST: data[CONF_HOST],
         CONF_TOKEN: data[CONF_TOKEN],
         CONF_MODEL: hub.device_info.model,
-        CONF_NAME: data[CONF_NAME],
+        CONF_NAME: name,
         CONF_MAC: format_mac(hub.device_info.mac_address),
     }
 
@@ -120,6 +130,7 @@ 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]
 
                 self.hass.config_entries.async_update_entry(existing_entry, data=data)
                 await self.hass.config_entries.async_reload(existing_entry.entry_id)
diff --git a/custom_components/xiaomi_viomi/strings.json b/custom_components/xiaomi_viomi/strings.json
index 561632d..9e00e43 100644
--- a/custom_components/xiaomi_viomi/strings.json
+++ b/custom_components/xiaomi_viomi/strings.json
@@ -3,6 +3,7 @@
     "error": {
       "cannot_connect": "[%key:common::config_flow::error::cannot_connect%]",
       "invalid_auth": "[%key:common::config_flow::error::invalid_auth%]",
+      "invalid_parameters": "[%key:common::config_flow::error::invalid_auth%]",
       "unknown": "[%key:common::config_flow::error::unknown%]"
     },
     "abort": {
@@ -13,10 +14,11 @@
       "user": {
         "data": {
           "host": "[%key:common::config_flow::data::host%]",
-          "token": "[%key:common::config_flow::data::api_token%]"
+          "token": "[%key:common::config_flow::data::api_token%]",
+          "name": "[%key:common::config_flow::data::name%]"
         },
         "description": "You will need the 32 character [%key:common::config_flow::data::api_token%], see https://www.home-assistant.io/integrations/xiaomi_miio#retrieving-the-access-token for instructions. Please note, that this [%key:common::config_flow::data::api_token%] is different from the key used by the Xiaomi Aqara integration.",
-        "title": "Connect to a Xiaomi Viomi Device2"
+        "title": "Connect to a Xiaomi Viomi Device"
       }
     }
   }
diff --git a/custom_components/xiaomi_viomi/translations/en.json b/custom_components/xiaomi_viomi/translations/en.json
index 2d2edea..beb3b83 100644
--- a/custom_components/xiaomi_viomi/translations/en.json
+++ b/custom_components/xiaomi_viomi/translations/en.json
@@ -6,6 +6,7 @@
     "error": {
       "invalid_auth": "Incomplete information to setup device, no host or token supplied",
       "cannot_connect": "Failed to connect",
+      "invalid_parameters": "Invalid parameters were provided",
       "unknown": "Unknown error occurred"
     },
     "flow_title": "{name}",
@@ -13,7 +14,8 @@
       "user": {
         "data": {
           "host": "IP Address",
-          "token": "API Token"
+          "token": "API Token",
+          "name": "Name of a device"
         },
         "description": "You will need the 32 character API Token, see https://www.home-assistant.io/integrations/xiaomi_miio#retrieving-the-access-token for instructions. Please note, that this API Token is different from the key used by the Xiaomi Aqara integration.",
         "title": "Connect to a Xiaomi Viomi Device"
diff --git a/custom_components/xiaomi_viomi/translations/ru.json b/custom_components/xiaomi_viomi/translations/ru.json
index aa77b63..dafe499 100644
--- a/custom_components/xiaomi_viomi/translations/ru.json
+++ b/custom_components/xiaomi_viomi/translations/ru.json
@@ -3,6 +3,7 @@
     "error": {
       "invalid_auth": "Неполная информация для настройки устройства, не указан хост или токен",
       "cannot_connect": "Не удалось подключиться",
+      "invalid_parameters": "Предоставлены неверные параметры устройства",
       "unknown": "Произошла неизвесная ошибка"
     },
     "abort": {
@@ -13,7 +14,8 @@
       "user": {
         "data": {
           "host": "IP-адрес",
-          "token": "Токен API"
+          "token": "Токен API",
+          "name": "Имя устройства"
         },
         "description": "Для подключения требуется 32-х значный Токен API. О том, как получить токен, Вы можете узнать здесь:\nhttps://www.home-assistant.io/integrations/xiaomi_miio#retrieving-the-access-token.\nОбратите внимание, что этот токен отличается от ключа, используемого при интеграции Xiaomi Aqara.",
         "title": "Подключение к устройству Xiaomi Viomi"
diff --git a/custom_components/xiaomi_viomi/translations/uk.json b/custom_components/xiaomi_viomi/translations/uk.json
index a2c83cf..d0e2b57 100644
--- a/custom_components/xiaomi_viomi/translations/uk.json
+++ b/custom_components/xiaomi_viomi/translations/uk.json
@@ -3,6 +3,7 @@
     "error": {
       "invalid_auth": "Неповна інформація для налагодження пристрою, не зазначений хост чи токен",
       "cannot_connect": "Не вдалося під'єднатися",
+      "invalid_parameters": "Надані невірні параметри пристрою",
       "unknown": "Трапилась невідома помилка"
     },
     "abort": {
@@ -13,7 +14,8 @@
       "user": {
         "data": {
           "host": "IP-адреса",
-          "token": "Токен API"
+          "token": "Токен API",
+          "name": "Ім'я пристрою"
         },
         "description": "Для підключення потрібно 32-х значний Токен API. Про те, як отримати токен, Ви можете дізнатися тут:\nhttps://www.home-assistant.io/integrations/vacuum.xiaomi_miio/#retrieving-the-access-token.\nЗверніть увагу, що цей токен відрізняється від ключа, який використовується при інтеграції Xiaomi Aqara.",
         "title": "Підключення до пристрою Xiaomi Viomi"
diff --git a/custom_components/xiaomi_viomi/vacuum.py b/custom_components/xiaomi_viomi/vacuum.py
index 782d5e0..b7ca4d4 100644
--- a/custom_components/xiaomi_viomi/vacuum.py
+++ b/custom_components/xiaomi_viomi/vacuum.py
@@ -58,7 +58,7 @@ async def async_setup_platform(
         domain=DOMAIN,
         data=config,
         version=2,
-        title=raw_config[CONF_NAME],
+        title=config[CONF_NAME],
         source=SOURCE_USER,
     )
     await async_setup_entry(hass, entry, async_add_entities)
@@ -188,7 +188,7 @@ def extra_state_attributes(self):
                     ATTR_FILTER_LEFT: int(
                         self.consumable_state.filter_left.total_seconds() / 3600
                     ),
-                    ATTR_STATUS: str(self._get_status()),
+                    ATTR_STATUS: self.state,
                     ATTR_MOP_ATTACHED: self.vacuum_state.mop_installed,
                 }
             )
@@ -211,10 +211,6 @@ def _got_error(self) -> bool:
         error_code = self.vacuum_state.error_code if self.vacuum_state else None
         return bool(error_code and error_code not in ERRORS_FALSE_POSITIVE)
 
-    def _get_status(self) -> str:
-        state_value = self.vacuum_state.state.value if self.vacuum_state else None
-        return STATE_CODE_TO_STATE[int(state_value)]
-
     def _get_device_status(self) -> ViomiVacuumStatus:
         """Override of miio's device.status() because of bug."""
         result = {}
diff --git a/tests/__init__.py b/tests/__init__.py
index f0491ec..695a93c 100644
--- a/tests/__init__.py
+++ b/tests/__init__.py
@@ -21,7 +21,6 @@
     "err_state": 2105,
     "has_map": 1,
     "has_newmap": 1,
-    "hw_info": "1.0.3",
     "is_charge": 0,
     "is_mop": 0,
     "is_work": 1,
@@ -36,7 +35,6 @@
     "s_time": 20,
     "start_time": 0,
     "suction_grade": 0,
-    "sw_info": "3.5.3_0017",
     "v_state": 10,
     "water_grade": 12,
     "zone_data": "0",
@@ -48,8 +46,8 @@
 }
 
 
-def get_entity_id() -> str:
-    return DOMAIN + "." + TEST_NAME
+def get_entity_id(use_model=False) -> str:
+    return DOMAIN + "." + (TEST_MODEL.replace(".", "_") if use_model else TEST_NAME)
 
 
 def get_mocked_entry():
@@ -66,13 +64,17 @@ def get_mocked_entry():
     )
 
 
-def mocked_viomi_device():
+def mocked_viomi_device(device_state_adjustment=None):
+    state = MOCKED_DEVICE_STATE
+    if device_state_adjustment is not None:
+        state = {**MOCKED_DEVICE_STATE, **device_state_adjustment}
+
     def _device_mock_method(command: str, parameters: Any = None):
         # Request for getting device state
         if command == "get_prop" and parameters:
             property_name = parameters[0]
-            if property_name in MOCKED_DEVICE_STATE:
-                return [MOCKED_DEVICE_STATE[property_name]]
+            if property_name in state:
+                return [state[property_name]]
 
             return [None]
         # Request for getting state of consumables
diff --git a/tests/test_config_flow.py b/tests/test_config_flow.py
index bda52e0..cd13789 100644
--- a/tests/test_config_flow.py
+++ b/tests/test_config_flow.py
@@ -4,9 +4,13 @@
 
 from homeassistant import config_entries, setup
 from homeassistant.core import HomeAssistant
-from homeassistant.data_entry_flow import RESULT_TYPE_CREATE_ENTRY, RESULT_TYPE_FORM
+from homeassistant.data_entry_flow import (
+    RESULT_TYPE_ABORT,
+    RESULT_TYPE_CREATE_ENTRY,
+    RESULT_TYPE_FORM,
+)
 
-from custom_components.xiaomi_viomi.config_flow import CannotConnect, InvalidAuth
+from custom_components.xiaomi_viomi.config_flow import CannotConnect
 from custom_components.xiaomi_viomi.const import DOMAIN
 from tests import TEST_HOST, TEST_MAC, TEST_MODEL, TEST_NAME, TEST_TOKEN
 
@@ -14,12 +18,12 @@
 async def test_form(hass: HomeAssistant) -> None:
     """Test we get the form."""
     await setup.async_setup_component(hass, "persistent_notification", {})
-    result = await hass.config_entries.flow.async_init(
+    flow_result = await hass.config_entries.flow.async_init(
         DOMAIN, context={"source": config_entries.SOURCE_USER}
     )
 
-    assert result["type"] == RESULT_TYPE_FORM
-    assert result["errors"] is None
+    assert flow_result["type"] == RESULT_TYPE_FORM
+    assert flow_result["errors"] is None
 
     device_info_mock = {"mac_address": TEST_MAC, "model": TEST_MODEL}
 
@@ -36,58 +40,79 @@ async def test_form(hass: HomeAssistant) -> None:
         "custom_components.xiaomi_viomi.async_setup_entry",
         return_value=True,
     ) as mock_setup_entry:
-        result2 = await hass.config_entries.flow.async_configure(
-            result["flow_id"],
+        result = await hass.config_entries.flow.async_configure(
+            flow_result["flow_id"],
             {
                 "host": TEST_HOST,
                 "token": TEST_TOKEN,
-                "model": TEST_MODEL,
                 "name": TEST_NAME,
             },
         )
         await hass.async_block_till_done()
 
-    result_data = {
-        "host": TEST_HOST,
-        "name": TEST_NAME,
-        "token": TEST_TOKEN,
-        "model": TEST_MODEL,
-        "mac": TEST_MAC,
-    }
+        result_data = {
+            "host": TEST_HOST,
+            "name": TEST_NAME,
+            "token": TEST_TOKEN,
+            "model": TEST_MODEL,
+            "mac": TEST_MAC,
+        }
 
-    assert result2["type"] == RESULT_TYPE_CREATE_ENTRY
-    assert result2["title"] == TEST_NAME
-    assert result2["data"] == result_data
-    assert len(mock_setup_entry.mock_calls) == 1
+        assert result["type"] == RESULT_TYPE_CREATE_ENTRY
+        assert result["title"] == TEST_NAME
+        assert result["data"] == result_data
+        assert len(mock_setup_entry.mock_calls) == 1
 
 
 async def test_form_invalid_auth(hass: HomeAssistant) -> None:
     """Test we handle invalid auth."""
-    result = await hass.config_entries.flow.async_init(
+    flow_result = await hass.config_entries.flow.async_init(
         DOMAIN, context={"source": config_entries.SOURCE_USER}
     )
 
     with patch(
         "custom_components.xiaomi_viomi.config_flow.ViomiDeviceHub.async_device_is_connectable",
-        side_effect=InvalidAuth,
+        return_value=False,
     ):
-        result2 = await hass.config_entries.flow.async_configure(
-            result["flow_id"],
+        result = await hass.config_entries.flow.async_configure(
+            flow_result["flow_id"],
             {
                 "host": TEST_HOST,
                 "token": TEST_TOKEN,
-                "model": TEST_MODEL,
                 "name": TEST_NAME,
             },
         )
 
-    assert result2["type"] == RESULT_TYPE_FORM
-    assert result2["errors"] == {"base": "invalid_auth"}
+        assert result["type"] == RESULT_TYPE_FORM
+        assert result["errors"] == {"base": "invalid_auth"}
+
+
+async def test_form_invalid_parameters(hass: HomeAssistant) -> None:
+    """Test we handle invalid auth."""
+    flow_result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": config_entries.SOURCE_USER}
+    )
+
+    with patch(
+        "custom_components.xiaomi_viomi.config_flow.ViomiDeviceHub.async_device_is_connectable",
+        side_effect=Exception,
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            flow_result["flow_id"],
+            {
+                "host": TEST_HOST,
+                "token": TEST_TOKEN,
+                "name": TEST_NAME,
+            },
+        )
+
+        assert result["type"] == RESULT_TYPE_FORM
+        assert result["errors"] == {"base": "unknown"}
 
 
 async def test_form_cannot_connect(hass: HomeAssistant) -> None:
     """Test we handle cannot connect error."""
-    result = await hass.config_entries.flow.async_init(
+    flow_result = await hass.config_entries.flow.async_init(
         DOMAIN, context={"source": config_entries.SOURCE_USER}
     )
 
@@ -95,15 +120,44 @@ async def test_form_cannot_connect(hass: HomeAssistant) -> None:
         "custom_components.xiaomi_viomi.config_flow.ViomiDeviceHub.async_device_is_connectable",
         side_effect=CannotConnect,
     ):
-        result2 = await hass.config_entries.flow.async_configure(
-            result["flow_id"],
+        result = await hass.config_entries.flow.async_configure(
+            flow_result["flow_id"],
             {
                 "host": TEST_HOST,
                 "token": TEST_TOKEN,
-                "model": TEST_MODEL,
                 "name": TEST_NAME,
             },
         )
 
-    assert result2["type"] == RESULT_TYPE_FORM
-    assert result2["errors"] == {"base": "cannot_connect"}
+        assert result["type"] == RESULT_TYPE_FORM
+        assert result["errors"] == {"base": "cannot_connect"}
+
+
+async def test_already_configured(hass: HomeAssistant) -> None:
+    """Test we handle cannot connect error."""
+    device_info_mock = {"mac_address": TEST_MAC, "model": TEST_MODEL}
+    with patch(
+        "custom_components.xiaomi_viomi.config_flow.ViomiDeviceHub.async_device_is_connectable",
+        return_value=True,
+    ), patch(
+        "custom_components.xiaomi_viomi.config_flow.ViomiDeviceHub.device_info",
+        new_callable=PropertyMock,
+        return_value=namedtuple("ObjectName", device_info_mock.keys())(
+            *device_info_mock.values()
+        ),
+    ):
+        for _ in range(0, 2):
+            flow_result = await hass.config_entries.flow.async_init(
+                DOMAIN, context={"source": config_entries.SOURCE_USER}
+            )
+            result = await hass.config_entries.flow.async_configure(
+                flow_result["flow_id"],
+                {
+                    "host": TEST_HOST,
+                    "token": TEST_TOKEN,
+                    "name": TEST_NAME,
+                },
+            )
+
+        assert result["type"] == RESULT_TYPE_ABORT
+        assert result["reason"] == "already_configured"
diff --git a/tests/test_setup.py b/tests/test_setup.py
index c6b8929..9b0bc03 100644
--- a/tests/test_setup.py
+++ b/tests/test_setup.py
@@ -21,7 +21,6 @@ async def test_platform_setup(hass: HomeAssistant):
                 "host": TEST_HOST,
                 "token": TEST_TOKEN,
                 "name": TEST_NAME,
-                "model": TEST_MODEL,
             }
         ]
     }
@@ -38,3 +37,30 @@ async def test_platform_setup(hass: HomeAssistant):
 
         assert state
         assert state.state == STATE_DOCKED
+        assert state.name == TEST_NAME
+
+
+async def test_platform_setup_without_name(hass: HomeAssistant):
+    config = {
+        DOMAIN: [
+            {
+                "platform": PLATFORM_NAME,
+                "host": TEST_HOST,
+                "token": TEST_TOKEN,
+            }
+        ]
+    }
+
+    with mocked_viomi_device() as mock_device_send:
+        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
diff --git a/tests/test_vacuum.py b/tests/test_vacuum.py
index 101f88c..25f5f6e 100644
--- a/tests/test_vacuum.py
+++ b/tests/test_vacuum.py
@@ -1,7 +1,10 @@
 """Test sensor for simple integration."""
+from typing import Optional
+
 import pytest
 from homeassistant.components.vacuum import (
     DOMAIN,
+    SERVICE_CLEAN_SPOT,
     SERVICE_LOCATE,
     SERVICE_PAUSE,
     SERVICE_RETURN_TO_BASE,
@@ -49,6 +52,7 @@ async def test_vacuum_state(hass: HomeAssistant):
     (SERVICE_RETURN_TO_BASE, "set_charge", [1]),
     (SERVICE_LOCATE, "set_resetpos", [1]),
     (SERVICE_STOP, "set_mode", [0, 0]),
+    (SERVICE_CLEAN_SPOT, None, None),
 ]
 
 
@@ -67,7 +71,9 @@ async def test_vacuum_service(hass: HomeAssistant, service, method, parameters):
         await hass.services.async_call(
             DOMAIN, service, {"entity_id": entity_id}, blocking=True
         )
-        mock_device_send.assert_any_call(method, parameters)
+
+        if method and parameters:
+            mock_device_send.assert_any_call(method, parameters)
 
 
 async def test_vacuum_send_command_service(hass: HomeAssistant):
@@ -92,15 +98,52 @@ async def test_vacuum_send_command_service(hass: HomeAssistant):
 
 
 @pytest.mark.parametrize(
-    "speed",
+    "service,initial_state,method,parameters",
     [
-        ViomiVacuumSpeed.Silent,
-        ViomiVacuumSpeed.Medium,
-        ViomiVacuumSpeed.Standard,
-        ViomiVacuumSpeed.Turbo,
+        (SERVICE_START_PAUSE, None, "set_mode_withroom", [0, 1, 0]),
+        (SERVICE_START_PAUSE, {"is_work": 0}, "set_mode", [0, 2]),
     ],
 )
-async def test_vacuum_fan_speed_service(hass: HomeAssistant, speed: ViomiVacuumSpeed):
+async def test_vacuum_start_stop(
+    hass: HomeAssistant, service, initial_state, method, parameters
+):
+    entry = get_mocked_entry()
+    with mocked_viomi_device(initial_state) as mock_device_send:
+        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()
+
+        await hass.services.async_call(
+            DOMAIN, service, {"entity_id": entity_id}, blocking=True
+        )
+
+        mock_device_send.assert_any_call(method, parameters)
+
+
+@pytest.mark.parametrize(
+    "speed,expected_value",
+    [
+        #  numeric values
+        (0, ViomiVacuumSpeed.Silent.value),
+        (1, ViomiVacuumSpeed.Standard.value),
+        (2, ViomiVacuumSpeed.Medium.value),
+        (3, ViomiVacuumSpeed.Turbo.value),
+        # string values
+        ("Silent", ViomiVacuumSpeed.Silent.value),
+        ("Standard", ViomiVacuumSpeed.Standard.value),
+        ("Medium", ViomiVacuumSpeed.Medium.value),
+        ("Turbo", ViomiVacuumSpeed.Turbo.value),
+        # Bad value
+        (-1, None),
+    ],
+)
+async def test_vacuum_fan_speed_service(
+    hass: HomeAssistant, speed: ViomiVacuumSpeed, expected_value: Optional[int]
+):
     entry = get_mocked_entry()
     with mocked_viomi_device() as mock_device_send:
         entry.add_to_hass(hass)
@@ -114,8 +157,59 @@ async def test_vacuum_fan_speed_service(hass: HomeAssistant, speed: ViomiVacuumS
         await hass.services.async_call(
             DOMAIN,
             SERVICE_SET_FAN_SPEED,
-            {"entity_id": entity_id, "fan_speed": speed.value},
+            {"entity_id": entity_id, "fan_speed": speed},
             blocking=True,
         )
 
-        mock_device_send.assert_any_call("set_suction", [speed.value])
+        if expected_value is not None:
+            mock_device_send.assert_any_call("set_suction", [expected_value])
+        else:
+            # Method wasn't called
+            with pytest.raises(AssertionError):
+                mock_device_send.assert_any_call("set_suction", [expected_value])
+
+
+@pytest.mark.parametrize(
+    "state_code,state_value",
+    [
+        (4, "returning"),
+        (9000, None),
+    ],
+)
+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:
+        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)
+
+        assert state
+        assert state.attributes["status"] == state_value
+
+
+@pytest.mark.parametrize(
+    "error_code,error_value",
+    [
+        (502, "Low battery"),
+        (9000, "Unknown error 9000"),
+    ],
+)
+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:
+        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)
+
+        assert state
+        assert state.attributes["error"] == error_value

From 14471c58f7d22ca53b77cb232937b18050e85c8e Mon Sep 17 00:00:00 2001
From: Artem Poliukhovych <nergal@users.noreply.github.com>
Date: Fri, 4 Feb 2022 13:59:44 +0200
Subject: [PATCH 4/8] Update python version for CI

---
 .github/workflows/workflow.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/workflow.yaml b/.github/workflows/workflow.yaml
index c65b352..897b6a1 100644
--- a/.github/workflows/workflow.yaml
+++ b/.github/workflows/workflow.yaml
@@ -13,7 +13,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: [3.8, 3.9, 3.10]
+        python-version: [3.8, 3.9, '3.10']
         os: [ubuntu-latest, macos-latest]
     runs-on: ${{ matrix.os }}
     steps:

From d6720d6536dd5ea5758c7080121f263348953e5a Mon Sep 17 00:00:00 2001
From: Artem Poliukhovych <nergal@users.noreply.github.com>
Date: Fri, 4 Feb 2022 14:02:47 +0200
Subject: [PATCH 5/8] Change coverage report format to xml

---
 Makefile | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Makefile b/Makefile
index 1eceb6c..c38592f 100644
--- a/Makefile
+++ b/Makefile
@@ -26,4 +26,4 @@ test:
 	$(POETRY) run pytest $(TEST_FOLDER)
 
 coverage:
-	$(POETRY) run pytest --cov-report term --cov=custom_components.xiaomi_viomi $(TEST_FOLDER)
\ No newline at end of file
+	$(POETRY) run pytest --cov-report xml --cov=custom_components.xiaomi_viomi $(TEST_FOLDER)
\ No newline at end of file

From 98b5f4497ce7d9c6128524ed50863966f3fc5c32 Mon Sep 17 00:00:00 2001
From: Artem Poliukhovych <nergal@users.noreply.github.com>
Date: Fri, 4 Feb 2022 14:08:07 +0200
Subject: [PATCH 6/8] Remove py 3.10 from the build because of the pip issue

---
 .github/workflows/workflow.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/workflow.yaml b/.github/workflows/workflow.yaml
index 897b6a1..b0105ac 100644
--- a/.github/workflows/workflow.yaml
+++ b/.github/workflows/workflow.yaml
@@ -13,7 +13,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: [3.8, 3.9, '3.10']
+        python-version: [3.8, 3.9]
         os: [ubuntu-latest, macos-latest]
     runs-on: ${{ matrix.os }}
     steps:

From 5725419a499faf98f43e44fa1fa44ca75cd62402 Mon Sep 17 00:00:00 2001
From: Artem Poliukhovych <nergal@users.noreply.github.com>
Date: Fri, 4 Feb 2022 14:10:07 +0200
Subject: [PATCH 7/8] Add codecov badge to readme

---
 README.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/README.md b/README.md
index 6a5b843..e3fecc6 100644
--- a/README.md
+++ b/README.md
@@ -1,6 +1,7 @@
 # Homeassistant Vacuum Viomi integration for HACS
 
 [![hacs_badge](https://img.shields.io/badge/HACS-Custom-orange.svg)](https://github.com/custom-components/hacs)
+[![codecov](https://codecov.io/gh/nergal/homeassistant-vacuum-viomi/branch/master/graph/badge.svg?token=MUHLPCJY2G)](https://codecov.io/gh/nergal/homeassistant-vacuum-viomi)
 [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
 ![GitHub](https://img.shields.io/github/license/nergal/homeassistant-vacuum-viomi)
 [![Validation flow](https://github.com/nergal/homeassistant-vacuum-viomi/actions/workflows/workflow.yaml/badge.svg)](https://github.com/nergal/homeassistant-vacuum-viomi/actions/workflows/workflow.yaml)

From 69de49d8327922f891f8c3ae3dbc630716a20a21 Mon Sep 17 00:00:00 2001
From: Artem Poliukhovych <nergal@users.noreply.github.com>
Date: Fri, 4 Feb 2022 14:15:09 +0200
Subject: [PATCH 8/8] Bump version to 0.0.4

---
 custom_components/xiaomi_viomi/manifest.json | 2 +-
 pyproject.toml                               | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/custom_components/xiaomi_viomi/manifest.json b/custom_components/xiaomi_viomi/manifest.json
index 11b8fec..5d65e1e 100644
--- a/custom_components/xiaomi_viomi/manifest.json
+++ b/custom_components/xiaomi_viomi/manifest.json
@@ -11,5 +11,5 @@
         "construct==2.10.67",
         "python-miio==0.5.9"
     ],
-    "version": "0.0.3"
+    "version": "0.0.4"
 }
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index 17bcdd5..e58f799 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "homeassistant-vacuum-viomi"
-version = "0.0.3"
+version = "0.0.4"
 description = "Integration of a Xiaomi Viomi vacuum devices"
 authors = ["Artem Poliukhovych <me@2mio.com>"]
 license = "MIT"