Merge pull request #103 from dmamontov/feature/vpn

Add support VPN

.gitignore

@@ -130,3 +130,7 @@ dmypy.json
 
 .DS_Store
 .idea
+.run
+bin
+share
+include

custom_components/miwifi/binary_sensor.py

@@ -22,6 +22,8 @@
     ATTR_BINARY_SENSOR_DUAL_BAND_NAME,
     ATTR_BINARY_SENSOR_WAN_STATE,
     ATTR_BINARY_SENSOR_WAN_STATE_NAME,
+    ATTR_BINARY_SENSOR_VPN_STATE,
+    ATTR_BINARY_SENSOR_VPN_STATE_NAME,
     ATTR_STATE,
     ATTR_STATE_NAME,
 )
@@ -51,6 +53,14 @@
         entity_category=EntityCategory.DIAGNOSTIC,
         entity_registry_enabled_default=True,
     ),
+    BinarySensorEntityDescription(
+        key=ATTR_BINARY_SENSOR_VPN_STATE,
+        name=ATTR_BINARY_SENSOR_VPN_STATE_NAME,
+        icon="mdi:security-network",
+        device_class=BinarySensorDeviceClass.CONNECTIVITY,
+        entity_category=EntityCategory.DIAGNOSTIC,
+        entity_registry_enabled_default=False,
+    ),
     BinarySensorEntityDescription(
         key=ATTR_BINARY_SENSOR_DUAL_BAND,
         name=ATTR_BINARY_SENSOR_DUAL_BAND_NAME,

custom_components/miwifi/const.py

@@ -127,6 +127,9 @@
 ATTR_SENSOR_UPTIME: Final = "uptime"
 ATTR_SENSOR_UPTIME_NAME: Final = "Uptime"
 
+ATTR_SENSOR_VPN_UPTIME: Final = "vpn_uptime"
+ATTR_SENSOR_VPN_UPTIME_NAME: Final = "Vpn uptime"
+
 ATTR_SENSOR_MEMORY_USAGE: Final = "memory_usage"
 ATTR_SENSOR_MEMORY_USAGE_NAME: Final = "Memory usage"
 
@@ -173,6 +176,9 @@
 ATTR_BINARY_SENSOR_DUAL_BAND: Final = "dual_band"
 ATTR_BINARY_SENSOR_DUAL_BAND_NAME: Final = "Dual band"
 
+ATTR_BINARY_SENSOR_VPN_STATE: Final = "vpn_state"
+ATTR_BINARY_SENSOR_VPN_STATE_NAME: Final = "Vpn state"
+
 """Light attributes"""
 ATTR_LIGHT_LED: Final = "led"
 ATTR_LIGHT_LED_NAME: Final = "Led"

custom_components/miwifi/luci.py

@@ -258,6 +258,14 @@ async def wifi_diag_detail_all(self) -> dict:
 
         return await self.get("xqnetwork/wifi_diag_detail_all")
 
+    async def vpn_status(self) -> dict:
+        """xqsystem/vpn_status method.
+
+        :return dict: dict with api data.
+        """
+
+        return await self.get("xqsystem/vpn_status")
+
     async def set_wifi(self, data: dict) -> dict:
         """xqnetwork/set_wifi method.
 

custom_components/miwifi/sensor.py

@@ -44,6 +44,8 @@
     ATTR_SENSOR_TEMPERATURE_NAME,
     ATTR_SENSOR_UPTIME,
     ATTR_SENSOR_UPTIME_NAME,
+    ATTR_SENSOR_VPN_UPTIME,
+    ATTR_SENSOR_VPN_UPTIME_NAME,
     ATTR_SENSOR_WAN_DOWNLOAD_SPEED,
     ATTR_SENSOR_WAN_DOWNLOAD_SPEED_NAME,
     ATTR_SENSOR_WAN_UPLOAD_SPEED,
@@ -77,6 +79,13 @@
         entity_category=EntityCategory.DIAGNOSTIC,
         entity_registry_enabled_default=False,
     ),
+    SensorEntityDescription(
+        key=ATTR_SENSOR_VPN_UPTIME,
+        name=ATTR_SENSOR_VPN_UPTIME_NAME,
+        icon="mdi:timer-sand",
+        entity_category=EntityCategory.DIAGNOSTIC,
+        entity_registry_enabled_default=False,
+    ),
     SensorEntityDescription(
         key=ATTR_SENSOR_MEMORY_USAGE,
         name=ATTR_SENSOR_MEMORY_USAGE_NAME,

custom_components/miwifi/updater.py

@@ -25,6 +25,7 @@
 
 from .const import (
     ATTR_BINARY_SENSOR_DUAL_BAND,
+    ATTR_BINARY_SENSOR_VPN_STATE,
     ATTR_BINARY_SENSOR_WAN_STATE,
     ATTR_DEVICE_HW_VERSION,
     ATTR_DEVICE_MAC_ADDRESS,
@@ -46,6 +47,7 @@
     ATTR_SENSOR_MODE,
     ATTR_SENSOR_TEMPERATURE,
     ATTR_SENSOR_UPTIME,
+    ATTR_SENSOR_VPN_UPTIME,
     ATTR_SENSOR_WAN_DOWNLOAD_SPEED,
     ATTR_SENSOR_WAN_UPLOAD_SPEED,
     ATTR_STATE,
@@ -102,6 +104,7 @@
 PREPARE_METHODS: Final = (
     "init",
     "status",
+    "vpn",
     "rom_update",
     "mode",
     "wan",
@@ -520,6 +523,27 @@ async def _async_prepare_status(self, data: dict) -> None:
             ) if "upspeed" in response["wan"] else 0
             # fmt: on
 
+    async def _async_prepare_vpn(self, data: dict) -> None:
+        """Prepare vpn.
+
+        :param data: dict
+        """
+
+        response: dict = await self.luci.vpn_status()
+
+        data |= {
+            ATTR_SENSOR_VPN_UPTIME: 0,
+            ATTR_BINARY_SENSOR_VPN_STATE: False,
+        }
+
+        if "uptime" in response:
+            data |= {
+                ATTR_SENSOR_VPN_UPTIME: str(
+                    timedelta(seconds=int(float(response["uptime"])))
+                ),
+                ATTR_BINARY_SENSOR_VPN_STATE: int(float(response["uptime"])) > 0,
+            }
+
     async def _async_prepare_rom_update(self, data: dict) -> None:
         """Prepare rom update.
 

pyvenv.cfg

@@ -0,0 +1,3 @@
+home = /opt/homebrew/opt/[email protected]/bin
+include-system-site-packages = false
+version = 3.10.8

tests/fixtures/vpn_status_data.json

@@ -0,0 +1,5 @@
+{
+  "status": 0,
+  "uptime": 343757,
+  "code": 0
+}

tests/fixtures/vpn_status_off_data.json

@@ -0,0 +1,5 @@
+{
+  "status": 4,
+  "uptime": 0,
+  "code": 0
+}

tests/setup.py

@@ -152,6 +152,9 @@ async def async_mock_luci_client(mock_luci_client) -> None:
     mock_luci_client.return_value.device_list = AsyncMock(
         return_value=json.loads(load_fixture("device_list_data.json"))
     )
+    mock_luci_client.return_value.vpn_status = AsyncMock(
+        return_value=json.loads(load_fixture("vpn_status_data.json"))
+    )
 
     async def mock_avaliable_channels(index: int = 1) -> dict:
         """Mock channels"""

tests/test_binary_sensors.py

@@ -26,6 +26,7 @@
 
 from custom_components.miwifi.const import (
     ATTR_BINARY_SENSOR_DUAL_BAND_NAME,
+    ATTR_BINARY_SENSOR_VPN_STATE_NAME,
     ATTR_BINARY_SENSOR_WAN_STATE_NAME,
     ATTR_DEVICE_MAC_ADDRESS,
     ATTR_STATE_NAME,
@@ -49,6 +50,7 @@ def auto_enable_custom_integrations(enable_custom_integrations):
     yield
 
 
+@pytest.mark.asyncio
 async def test_init(hass: HomeAssistant) -> None:
     """Test init.
 
@@ -100,7 +102,13 @@ async def test_init(hass: HomeAssistant) -> None:
         )
         assert dual_band_state is None
 
+        vpn_state: State | None = hass.states.get(
+            _generate_id(ATTR_BINARY_SENSOR_VPN_STATE_NAME, updater)
+        )
+        assert vpn_state is None
+
 
+@pytest.mark.asyncio
 async def test_update_state(hass: HomeAssistant) -> None:
     """Test update state.
 
@@ -172,6 +180,7 @@ def error() -> None:
         assert state.attributes["icon"] == "mdi:router-wireless-off"
 
 
+@pytest.mark.asyncio
 async def test_update_wan_state(hass: HomeAssistant) -> None:
     """Test update wan state.
 
@@ -251,6 +260,7 @@ def _off() -> dict:
         assert state.state == STATE_UNAVAILABLE
 
 
+@pytest.mark.asyncio
 async def test_update_dual_band(hass: HomeAssistant) -> None:
     """Test update dual_band.
 
@@ -350,6 +360,106 @@ def _on() -> dict:
         assert state.state == STATE_UNAVAILABLE
 
 
+@pytest.mark.asyncio
+async def test_update_vpn_state(hass: HomeAssistant) -> None:
+    """Test update vpn_state.
+
+    :param hass: HomeAssistant
+    """
+
+    with patch(
+        "custom_components.miwifi.updater.LuciClient"
+    ) as mock_luci_client, patch(
+        "custom_components.miwifi.async_start_discovery", return_value=None
+    ), patch(
+        "custom_components.miwifi.device_tracker.socket.socket"
+    ) as mock_socket, patch(
+        "custom_components.miwifi.updater.asyncio.sleep", return_value=None
+    ):
+        await async_mock_luci_client(mock_luci_client)
+
+        mock_socket.return_value.recv.return_value = AsyncMock(return_value=None)
+
+        def success() -> dict:
+            return json.loads(load_fixture("device_list_data.json"))
+
+        def error() -> None:
+            raise LuciRequestError
+
+        mock_luci_client.return_value.device_list = AsyncMock(
+            side_effect=MultipleSideEffect(success, success, success, error, error)
+        )
+
+        def _off() -> dict:
+            return json.loads(load_fixture("vpn_status_off_data.json"))
+
+        def _on() -> dict:
+            return json.loads(load_fixture("vpn_status_data.json"))
+
+        mock_luci_client.return_value.vpn_status = AsyncMock(
+            side_effect=MultipleSideEffect(_off, _off, _off, _on, _on)
+        )
+
+        setup_data: list = await async_setup(hass)
+
+        config_entry: MockConfigEntry = setup_data[1]
+
+        assert await hass.config_entries.async_setup(config_entry.entry_id)
+        await hass.async_block_till_done()
+
+        updater: LuciUpdater = hass.data[DOMAIN][config_entry.entry_id][UPDATER]
+        registry = er.async_get(hass)
+
+        assert updater.last_update_success
+
+        unique_id: str = _generate_id(ATTR_BINARY_SENSOR_VPN_STATE_NAME, updater)
+
+        entry: er.RegistryEntry | None = registry.async_get(unique_id)
+        state: State = hass.states.get(unique_id)
+        assert state is None
+        assert entry is not None
+        assert entry.disabled_by == er.RegistryEntryDisabler.INTEGRATION
+
+        registry.async_update_entity(entity_id=unique_id, disabled_by=None)
+        await hass.async_block_till_done()
+
+        async_fire_time_changed(
+            hass, utcnow() + timedelta(seconds=DEFAULT_SCAN_INTERVAL + 1)
+        )
+        await hass.async_block_till_done()
+
+        async_fire_time_changed(
+            hass, utcnow() + timedelta(seconds=DEFAULT_SCAN_INTERVAL + 1)
+        )
+        await hass.async_block_till_done()
+
+        entry = registry.async_get(unique_id)
+        state = hass.states.get(unique_id)
+
+        assert entry.disabled_by is None
+        assert state.state == STATE_OFF
+        assert state.name == ATTR_BINARY_SENSOR_VPN_STATE_NAME
+        assert state.attributes["icon"] == "mdi:security-network"
+        assert state.attributes["attribution"] == ATTRIBUTION
+        assert entry.entity_category == EntityCategory.DIAGNOSTIC
+
+        async_fire_time_changed(
+            hass, utcnow() + timedelta(seconds=DEFAULT_SCAN_INTERVAL + 1)
+        )
+        await hass.async_block_till_done()
+
+        state = hass.states.get(unique_id)
+        assert state.state == STATE_ON
+
+        async_fire_time_changed(
+            hass, utcnow() + timedelta(seconds=DEFAULT_SCAN_INTERVAL + 1)
+        )
+        await hass.async_block_till_done()
+
+        state = hass.states.get(unique_id)
+        assert state.state == STATE_UNAVAILABLE
+
+
 def _generate_id(code: str, updater: LuciUpdater) -> str:
     """Generate unique id
 

tests/test_button.py

@@ -47,6 +47,7 @@ def auto_enable_custom_integrations(enable_custom_integrations):
     yield
 
 
+@pytest.mark.asyncio
 async def test_init(hass: HomeAssistant) -> None:
     """Test init.
 
@@ -85,6 +86,7 @@ async def test_init(hass: HomeAssistant) -> None:
         assert state.attributes["attribution"] == ATTRIBUTION
 
 
+@pytest.mark.asyncio
 async def test_update_reboot(hass: HomeAssistant) -> None:
     """Test update reboot.