Improve type hints in mfi (home-assistant#160985)

Author: epenet

homeassistant/components/mfi/sensor.py

@@ -4,14 +4,15 @@
 
 import logging
 
-from mficlient.client import FailedToLogin, MFiClient
+from mficlient.client import FailedToLogin, MFiClient, Port as MFiPort
 import requests
 import voluptuous as vol
 
 from homeassistant.components.sensor import (
     PLATFORM_SCHEMA as SENSOR_PLATFORM_SCHEMA,
     SensorDeviceClass,
     SensorEntity,
+    StateType,
 )
 from homeassistant.const import (
     CONF_HOST,
@@ -64,24 +65,29 @@ def setup_platform(
     discovery_info: DiscoveryInfoType | None = None,
 ) -> None:
     """Set up mFi sensors."""
-    host = config.get(CONF_HOST)
-    username = config.get(CONF_USERNAME)
-    password = config.get(CONF_PASSWORD)
-    use_tls = config.get(CONF_SSL)
-    verify_tls = config.get(CONF_VERIFY_SSL)
+    host: str = config[CONF_HOST]
+    username: str = config[CONF_USERNAME]
+    password: str = config[CONF_PASSWORD]
+    use_tls: bool = config[CONF_SSL]
+    verify_tls: bool = config[CONF_VERIFY_SSL]
     default_port = 6443 if use_tls else 6080
-    port = int(config.get(CONF_PORT, default_port))
+    network_port: int = config.get(CONF_PORT, default_port)
 
     try:
         client = MFiClient(
-            host, username, password, port=port, use_tls=use_tls, verify=verify_tls
+            host,
+            username,
+            password,
+            port=network_port,
+            use_tls=use_tls,
+            verify=verify_tls,
         )
     except (FailedToLogin, requests.exceptions.ConnectionError) as ex:
         _LOGGER.error("Unable to connect to mFi: %s", str(ex))
         return
 
     add_entities(
-        MfiSensor(port, hass)
+        MfiSensor(port)
         for device in client.get_devices()
         for port in device.ports.values()
         if port.model in SENSOR_MODELS
@@ -91,18 +97,17 @@ def setup_platform(
 class MfiSensor(SensorEntity):
     """Representation of a mFi sensor."""
 
-    def __init__(self, port, hass):
+    def __init__(self, port: MFiPort) -> None:
         """Initialize the sensor."""
         self._port = port
-        self._hass = hass
 
     @property
-    def name(self):
+    def name(self) -> str:
         """Return the name of the sensor."""
         return self._port.label
 
     @property
-    def native_value(self):
+    def native_value(self) -> StateType:
         """Return the state of the sensor."""
         try:
             tag = self._port.tag
@@ -129,7 +134,7 @@ def device_class(self) -> SensorDeviceClass | None:
         return None
 
     @property
-    def native_unit_of_measurement(self):
+    def native_unit_of_measurement(self) -> str | None:
         """Return the unit of measurement of this entity, if any."""
         try:
             tag = self._port.tag

homeassistant/components/mfi/switch.py

@@ -5,7 +5,7 @@
 import logging
 from typing import Any
 
-from mficlient.client import FailedToLogin, MFiClient
+from mficlient.client import FailedToLogin, MFiClient, Port as MFiPort
 import requests
 import voluptuous as vol
 
@@ -51,18 +51,23 @@ def setup_platform(
     add_entities: AddEntitiesCallback,
     discovery_info: DiscoveryInfoType | None = None,
 ) -> None:
-    """Set up mFi sensors."""
-    host = config.get(CONF_HOST)
-    username = config.get(CONF_USERNAME)
-    password = config.get(CONF_PASSWORD)
-    use_tls = config[CONF_SSL]
-    verify_tls = config.get(CONF_VERIFY_SSL)
+    """Set up mFi switches."""
+    host: str = config[CONF_HOST]
+    username: str = config[CONF_USERNAME]
+    password: str = config[CONF_PASSWORD]
+    use_tls: bool = config[CONF_SSL]
+    verify_tls: bool = config[CONF_VERIFY_SSL]
     default_port = 6443 if use_tls else 6080
-    port = int(config.get(CONF_PORT, default_port))
+    network_port: int = config.get(CONF_PORT, default_port)
 
     try:
         client = MFiClient(
-            host, username, password, port=port, use_tls=use_tls, verify=verify_tls
+            host,
+            username,
+            password,
+            port=network_port,
+            use_tls=use_tls,
+            verify=verify_tls,
         )
     except (FailedToLogin, requests.exceptions.ConnectionError) as ex:
         _LOGGER.error("Unable to connect to mFi: %s", str(ex))
@@ -79,23 +84,23 @@ def setup_platform(
 class MfiSwitch(SwitchEntity):
     """Representation of an mFi switch-able device."""
 
-    def __init__(self, port):
+    def __init__(self, port: MFiPort) -> None:
         """Initialize the mFi device."""
         self._port = port
-        self._target_state = None
+        self._target_state: bool | None = None
 
     @property
-    def unique_id(self):
+    def unique_id(self) -> str:
         """Return the unique ID of the device."""
         return self._port.ident
 
     @property
-    def name(self):
+    def name(self) -> str:
         """Return the name of the device."""
         return self._port.label
 
     @property
-    def is_on(self):
+    def is_on(self) -> bool | None:
         """Return true if the device is on."""
         return self._port.output
 

tests/components/mfi/test_sensor.py

@@ -34,22 +34,22 @@ async def test_setup_missing_config(hass: HomeAssistant) -> None:
     """Test setup with missing configuration."""
     with mock.patch("homeassistant.components.mfi.sensor.MFiClient") as mock_client:
         config = {"sensor": {"platform": "mfi"}}
-        assert await async_setup_component(hass, "sensor", config)
+        assert await async_setup_component(hass, COMPONENT.DOMAIN, config)
         assert not mock_client.called
 
 
 async def test_setup_failed_login(hass: HomeAssistant) -> None:
     """Test setup with login failure."""
     with mock.patch("homeassistant.components.mfi.sensor.MFiClient") as mock_client:
         mock_client.side_effect = FailedToLogin
-        assert not PLATFORM.setup_platform(hass, GOOD_CONFIG, None)
+        assert not PLATFORM.setup_platform(hass, GOOD_CONFIG[COMPONENT.DOMAIN], None)
 
 
 async def test_setup_failed_connect(hass: HomeAssistant) -> None:
     """Test setup with connection failure."""
     with mock.patch("homeassistant.components.mfi.sensor.MFiClient") as mock_client:
         mock_client.side_effect = requests.exceptions.ConnectionError
-        assert not PLATFORM.setup_platform(hass, GOOD_CONFIG, None)
+        assert not PLATFORM.setup_platform(hass, GOOD_CONFIG[COMPONENT.DOMAIN], None)
 
 
 async def test_setup_minimum(hass: HomeAssistant) -> None:
@@ -111,7 +111,7 @@ async def test_setup_adds_proper_devices(hass: HomeAssistant) -> None:
         await hass.async_block_till_done()
         for ident, port in ports.items():
             if ident != "bad":
-                mock_sensor.assert_any_call(port, hass)
+                mock_sensor.assert_any_call(port)
         assert mock.call(ports["bad"], hass) not in mock_sensor.mock_calls
 
 
@@ -124,7 +124,7 @@ def port_fixture() -> mock.MagicMock:
 @pytest.fixture(name="sensor")
 def sensor_fixture(hass: HomeAssistant, port: mock.MagicMock) -> mfi.MfiSensor:
     """Sensor fixture."""
-    sensor = mfi.MfiSensor(port, hass)
+    sensor = mfi.MfiSensor(port)
     sensor.hass = hass
     return sensor