Refine code

custom_components/jq300/__init__.py

@@ -20,6 +20,7 @@
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
+from homeassistant.helpers.typing import ConfigType
 
 from .api import Jq300Account
 from .const import (
@@ -52,16 +53,17 @@
 CONFIG_SCHEMA = vol.Schema({DOMAIN: ACCOUNT_SCHEMA}, extra=vol.ALLOW_EXTRA)
 
 
-async def async_setup(hass: HomeAssistant, config):
+async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     """Set up this integration using YAML."""
+    # Print startup message
+    if DOMAIN not in hass.data:
+        _LOGGER.info(STARTUP_MESSAGE)
+        hass.data[DOMAIN] = {}
+
     if DOMAIN not in config:
         return True
 
-    # Print startup message
-    _LOGGER.info(STARTUP_MESSAGE)
-    hass.data.setdefault(DOMAIN, {})
-
-    hass.data[DOMAIN][CONF_YAML] = config.get(DOMAIN)
+    hass.data[DOMAIN][CONF_YAML] = config[DOMAIN]
     hass.async_create_task(
         hass.config_entries.flow.async_init(
             DOMAIN, context={"source": SOURCE_IMPORT}, data={}

custom_components/jq300/binary_sensor.py

@@ -13,7 +13,9 @@
 import logging
 
 from homeassistant.components.binary_sensor import ENTITY_ID_FORMAT, BinarySensorEntity
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_DEVICE_CLASS, CONF_DEVICES, CONF_ICON, CONF_NAME
+from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity import async_generate_entity_id
 
 from .api import Jq300Account
@@ -23,7 +25,9 @@
 _LOGGER = logging.getLogger(__name__)
 
 
-async def async_setup_entry(hass, entry, async_add_entities):
+async def async_setup_entry(
+    hass: HomeAssistant, entry: ConfigEntry, async_add_entities
+) -> bool:
     """Set up binary_sensor platform."""
     data = hass.data[DOMAIN][entry.entry_id]
     account = data[CONF_ACCOUNT_CONTROLLER]  # type: Jq300Account

custom_components/jq300/entity.py

@@ -11,20 +11,12 @@
 
 import logging
 
-from homeassistant.const import ATTR_ATTRIBUTION, ATTR_DEVICE_ID
+from homeassistant.const import ATTR_ATTRIBUTION
 from homeassistant.exceptions import PlatformNotReady
 from homeassistant.helpers.entity import Entity
 
-from custom_components.jq300 import Jq300Account
-
-from .const import (
-    ATTR_DEVICE_BRAND,
-    ATTR_DEVICE_MODEL,
-    ATTRIBUTION,
-    DOMAIN,
-    NAME,
-    VERSION,
-)
+from .api import Jq300Account
+from .const import ATTRIBUTION, DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -91,16 +83,14 @@ def device_info(self):
         """Return the device info."""
         return {
             "identifiers": {(DOMAIN, self.unique_id)},
-            "name": NAME,
-            "model": VERSION,
+            "name": self._name,
+            "manufacturer": self._device.get("brandname"),
+            "model": self._device.get("pt_model"),
         }
 
     @property
     def device_state_attributes(self):
         """Return the state attributes."""
         return {
-            ATTR_DEVICE_BRAND: self._device.get("brandname"),
-            ATTR_DEVICE_MODEL: self._device.get("pt_model"),
-            ATTR_DEVICE_ID: self._device.get("deviceid"),
             ATTR_ATTRIBUTION: ATTRIBUTION,
         }

tests/test_binary_sensor.py

@@ -6,7 +6,7 @@
 
 
 async def test_entity_initialization(mock_account: Jq300Account):
-    """Test sensor initialization."""
+    """Test entity initialization."""
     mock_account._devices = {123: {"pt_name": "Kitchen"}}
 
     entity = Jq300BinarySensor("test", mock_account, 123, 1, False)

tests/test_config_flow.py

@@ -1,9 +1,12 @@
-"""Test integration_blueprint config flow."""
+"""Test jq300 config flow."""
+
+from unittest.mock import patch
 
 from homeassistant import config_entries, data_entry_flow
 from homeassistant.core import HomeAssistant
 
 from custom_components.jq300 import DOMAIN
+from custom_components.jq300.config_flow import Jq300FlowHandler
 
 
 async def test_async_step_import(hass: HomeAssistant):
@@ -17,3 +20,14 @@ async def test_async_step_import(hass: HomeAssistant):
 
     assert result["type"] == data_entry_flow.RESULT_TYPE_CREATE_ENTRY
     assert result["data"] == config
+
+    with patch.object(
+        Jq300FlowHandler, "_async_current_entries", return_value=["Test"]
+    ):
+        result = await hass.config_entries.flow.async_init(
+            DOMAIN,
+            context={"source": config_entries.SOURCE_IMPORT},
+            data=config,
+        )
+
+        assert result["type"] == data_entry_flow.RESULT_TYPE_ABORT

tests/test_entity.py

@@ -5,29 +5,29 @@
 from pytest import raises
 
 from custom_components.jq300 import Jq300Account
-from custom_components.jq300.const import ATTRIBUTION, DOMAIN, NAME, VERSION
+from custom_components.jq300.const import ATTRIBUTION, DOMAIN
 from custom_components.jq300.entity import Jq300Entity
 
 
 async def test_entity_initialization(mock_account: Jq300Account):
-    """Test sensor initialization."""
+    """Test entity initialization."""
     with raises(PlatformNotReady):
         _ = Jq300Entity("test", mock_account, 123, 7, 12)
 
-    mock_account._devices = {123: {"pt_name": "Kitchen"}}
+    mock_account._devices = {
+        123: {"pt_name": "Kitchen", "brandname": "qwe", "pt_model": "asd"}
+    }
 
     entity = Jq300Entity("test", mock_account, 123, 7, 12)
 
     expected_device_info = {
         "identifiers": {(DOMAIN, "test@email.com-123-7")},
-        "model": VERSION,
-        "name": NAME,
+        "name": None,
+        "manufacturer": "qwe",
+        "model": "asd",
     }
     expected_attributes = {
         "attribution": ATTRIBUTION,
-        "device_brand": None,
-        "device_id": None,
-        "device_model": None,
     }
 
     assert entity.unique_id == "test@email.com-123-7"

tests/test_sensor.py

@@ -9,16 +9,13 @@
 
 
 async def test_entity_initialization(mock_account: Jq300Account):
-    """Test sensor initialization."""
+    """Test entity initialization."""
     mock_account._devices = {123: {"pt_name": "Kitchen"}}
 
     entity = Jq300Sensor("test", mock_account, 123, 7, 12)
 
     expected_attributes = {
         "attribution": ATTRIBUTION,
-        "device_brand": None,
-        "device_id": None,
-        "device_model": None,
         "raw_state": 12,
     }