Add WAN port connection control to TP-Link Omada gateways (#111431)

* Add gateway connection switches

* Add unit tests

* tplink_omada code review feedback

* Rebase and move icons and strings into json

homeassistant/components/tplink_omada/__init__.py

@@ -47,6 +47,10 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     site_client = await client.get_site_client(OmadaSite("", entry.data[CONF_SITE]))
     controller = OmadaSiteController(hass, site_client)
+    gateway_coordinator = await controller.get_gateway_coordinator()
+    if gateway_coordinator:
+        await gateway_coordinator.async_config_entry_first_refresh()
+
     hass.data[DOMAIN][entry.entry_id] = controller
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)

homeassistant/components/tplink_omada/binary_sensor.py

@@ -2,19 +2,21 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable, Generator
+from collections.abc import Callable
+from dataclasses import dataclass
 
-from attr import dataclass
+from tplink_omada_client.definitions import GatewayPortMode, LinkStatus, PoEMode
-from tplink_omada_client.definitions import GatewayPortMode, LinkStatus
 from tplink_omada_client.devices import (
     OmadaDevice,
     OmadaGateway,
+    OmadaGatewayPortConfig,
     OmadaGatewayPortStatus,
 )
 
 from homeassistant.components.binary_sensor import (
     BinarySensorDeviceClass,
     BinarySensorEntity,
+    BinarySensorEntityDescription,
 )
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
@@ -32,94 +34,103 @@ async def async_setup_entry(
 ) -> None:
     """Set up binary sensors."""
     controller: OmadaSiteController = hass.data[DOMAIN][config_entry.entry_id]
-    omada_client = controller.omada_client
 
     gateway_coordinator = await controller.get_gateway_coordinator()
     if not gateway_coordinator:
         return
 
-    gateway = await omada_client.get_gateway(gateway_coordinator.mac)
+    entities: list[OmadaDeviceEntity] = []
-
+    for gateway in gateway_coordinator.data.values():
-    async_add_entities(
+        entities.extend(
-        get_gateway_port_status_sensors(gateway, hass, gateway_coordinator)
+            OmadaGatewayPortBinarySensor(
-    )
+                gateway_coordinator, gateway, p.port_number, desc
-
-    await gateway_coordinator.async_request_refresh()
-
-
-def get_gateway_port_status_sensors(
-    gateway: OmadaGateway, hass: HomeAssistant, coordinator: OmadaGatewayCoordinator
-) -> Generator[BinarySensorEntity, None, None]:
-    """Generate binary sensors for gateway ports."""
-    for port in gateway.port_status:
-        if port.mode == GatewayPortMode.WAN:
-            yield OmadaGatewayPortBinarySensor(
-                coordinator,
-                gateway,
-                GatewayPortBinarySensorConfig(
-                    port_number=port.port_number,
-                    id_suffix="wan_link",
-                    name_suffix="Internet Link",
-                    device_class=BinarySensorDeviceClass.CONNECTIVITY,
-                    update_func=lambda p: p.wan_connected,
-                ),
-            )
-        if port.mode == GatewayPortMode.LAN:
-            yield OmadaGatewayPortBinarySensor(
-                coordinator,
-                gateway,
-                GatewayPortBinarySensorConfig(
-                    port_number=port.port_number,
-                    id_suffix="lan_status",
-                    name_suffix="LAN Status",
-                    device_class=BinarySensorDeviceClass.CONNECTIVITY,
-                    update_func=lambda p: p.link_status == LinkStatus.LINK_UP,
-                ),
             )
+            for p in gateway.port_configs
+            for desc in GATEWAY_PORT_SENSORS
+            if desc.exists_func(p)
+        )
+
+    async_add_entities(entities)
 
 
-@dataclass
+@dataclass(frozen=True, kw_only=True)
-class GatewayPortBinarySensorConfig:
+class GatewayPortBinarySensorEntityDescription(BinarySensorEntityDescription):
-    """Config for a binary status derived from a gateway port."""
+    """Entity description for a binary status derived from a gateway port."""
 
-    port_number: int
+    exists_func: Callable[[OmadaGatewayPortConfig], bool] = lambda _: True
-    id_suffix: str
-    name_suffix: str
-    device_class: BinarySensorDeviceClass
     update_func: Callable[[OmadaGatewayPortStatus], bool]
 
 
+GATEWAY_PORT_SENSORS: list[GatewayPortBinarySensorEntityDescription] = [
+    GatewayPortBinarySensorEntityDescription(
+        key="wan_link",
+        translation_key="wan_link",
+        device_class=BinarySensorDeviceClass.CONNECTIVITY,
+        exists_func=lambda p: p.port_status.mode == GatewayPortMode.WAN,
+        update_func=lambda p: p.wan_connected,
+    ),
+    GatewayPortBinarySensorEntityDescription(
+        key="online_detection",
+        translation_key="online_detection",
+        device_class=BinarySensorDeviceClass.CONNECTIVITY,
+        exists_func=lambda p: p.port_status.mode == GatewayPortMode.WAN,
+        update_func=lambda p: p.online_detection,
+    ),
+    GatewayPortBinarySensorEntityDescription(
+        key="lan_status",
+        translation_key="lan_status",
+        device_class=BinarySensorDeviceClass.CONNECTIVITY,
+        exists_func=lambda p: p.port_status.mode == GatewayPortMode.LAN,
+        update_func=lambda p: p.link_status == LinkStatus.LINK_UP,
+    ),
+    GatewayPortBinarySensorEntityDescription(
+        key="poe_delivery",
+        translation_key="poe_delivery",
+        device_class=BinarySensorDeviceClass.POWER,
+        exists_func=lambda p: (
+            p.port_status.mode == GatewayPortMode.LAN and p.poe_mode == PoEMode.ENABLED
+        ),
+        update_func=lambda p: p.poe_active,
+    ),
+]
+
+
 class OmadaGatewayPortBinarySensor(OmadaDeviceEntity[OmadaGateway], BinarySensorEntity):
     """Binary status of a property on an internet gateway."""
 
+    entity_description: GatewayPortBinarySensorEntityDescription
     _attr_has_entity_name = True
 
     def __init__(
         self,
         coordinator: OmadaGatewayCoordinator,
         device: OmadaDevice,
-        config: GatewayPortBinarySensorConfig,
+        port_number: int,
+        entity_description: GatewayPortBinarySensorEntityDescription,
     ) -> None:
         """Initialize the gateway port binary sensor."""
         super().__init__(coordinator, device)
-        self._config = config
+        self.entity_description = entity_description
-        self._attr_unique_id = f"{device.mac}_{config.port_number}_{config.id_suffix}"
+        self._port_number = port_number
-        self._attr_device_class = config.device_class
+        self._attr_unique_id = f"{device.mac}_{port_number}_{entity_description.key}"
+        self._attr_translation_placeholders = {"port_name": f"{port_number}"}
 
-        self._attr_name = f"Port {config.port_number} {config.name_suffix}"
+    async def async_added_to_hass(self) -> None:
+        """When entity is added to hass."""
+        await super().async_added_to_hass()
+        self._do_update()
+
+    def _do_update(self) -> None:
+        gateway = self.coordinator.data[self.device.mac]
+
+        port = next(
+            p for p in gateway.port_status if p.port_number == self._port_number
+        )
+        if port:
+            self._attr_is_on = self.entity_description.update_func(port)
 
     @callback
     def _handle_coordinator_update(self) -> None:
         """Handle updated data from the coordinator."""
-        gateway = self.coordinator.data[self.device.mac]
+        self._do_update()
-
-        port = next(
-            p for p in gateway.port_status if p.port_number == self._config.port_number
-        )
-        if port:
-            self._attr_is_on = self._config.update_func(port)
-            self._attr_available = True
-        else:
-            self._attr_available = False
-
         self.async_write_ha_state()

homeassistant/components/tplink_omada/icons.json

@@ -3,6 +3,20 @@
     "switch": {
       "poe_control": {
         "default": "mdi:ethernet"
+      },
+      "wan_connect_ipv4": {
+        "default": "mdi:wan"
+      },
+      "wan_connect_ipv6": {
+        "default": "mdi:wan"
+      }
+    },
+    "binary_sensor": {
+      "online_detection": {
+        "default": "mdi:cloud-check",
+        "state": {
+          "off": "mdi:cloud-cancel"
+        }
       }
     }
   }

homeassistant/components/tplink_omada/strings.json

@@ -39,5 +39,32 @@
       "already_configured": "[%key:common::config_flow::abort::already_configured_device%]",
       "reauth_successful": "[%key:common::config_flow::abort::reauth_successful%]"
     }
+  },
+  "entity": {
+    "switch": {
+      "poe_control": {
+        "name": "Port {port_name} PoE"
+      },
+      "wan_connect_ipv4": {
+        "name": "Port {port_name} Internet Connected"
+      },
+      "wan_connect_ipv6": {
+        "name": "Port {port_name} Internet Connected (IPv6)"
+      }
+    },
+    "binary_sensor": {
+      "wan_link": {
+        "name": "Port {port_name} Internet Link"
+      },
+      "online_detection": {
+        "name": "Port {port_name} Online Detection"
+      },
+      "lan_status": {
+        "name": "Port {port_name} LAN Status"
+      },
+      "poe_delivery": {
+        "name": "Port {port_name} PoE Delivery"
+      }
+    }
   }
 }

homeassistant/components/tplink_omada/switch.py

@@ -2,20 +2,33 @@
 
 from __future__ import annotations
 
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
+from functools import partial
 from typing import Any
 
-from tplink_omada_client import SwitchPortOverrides
+from tplink_omada_client import OmadaSiteClient, SwitchPortOverrides
-from tplink_omada_client.definitions import PoEMode
+from tplink_omada_client.definitions import GatewayPortMode, PoEMode
-from tplink_omada_client.devices import OmadaSwitch, OmadaSwitchPortDetails
+from tplink_omada_client.devices import (
+    OmadaDevice,
+    OmadaGateway,
+    OmadaGatewayPortStatus,
+    OmadaSwitch,
+    OmadaSwitchPortDetails,
+)
 
-from homeassistant.components.switch import SwitchEntity
+from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from .const import DOMAIN
-from .controller import OmadaSiteController, OmadaSwitchPortCoordinator
+from .controller import (
+    OmadaGatewayCoordinator,
+    OmadaSiteController,
+    OmadaSwitchPortCoordinator,
+)
 from .entity import OmadaDeviceEntity
 
 
@@ -44,15 +57,69 @@ async def async_setup_entry(
                     OmadaNetworkSwitchPortPoEControl(coordinator, switch, port_id)
                 )
 
+    gateway_coordinator = await controller.get_gateway_coordinator()
+    if gateway_coordinator:
+        for gateway in gateway_coordinator.data.values():
+            entities.extend(
+                OmadaGatewayPortSwitchEntity(
+                    gateway_coordinator, gateway, p.port_number, desc
+                )
+                for p in gateway.port_status
+                for desc in GATEWAY_PORT_SWITCHES
+                if desc.exists_func(p)
+            )
+
     async_add_entities(entities)
 
 
+@dataclass(frozen=True, kw_only=True)
+class GatewayPortSwitchEntityDescription(SwitchEntityDescription):
+    """Entity description for a toggle switch derived from a gateway port."""
+
+    exists_func: Callable[[OmadaGatewayPortStatus], bool] = lambda _: True
+    set_func: Callable[
+        [OmadaSiteClient, OmadaDevice, OmadaGatewayPortStatus, bool],
+        Awaitable[OmadaGatewayPortStatus],
+    ]
+    update_func: Callable[[OmadaGatewayPortStatus], bool]
+
+
+def _wan_connect_disconnect(
+    client: OmadaSiteClient,
+    device: OmadaDevice,
+    port: OmadaGatewayPortStatus,
+    enable: bool,
+    ipv6: bool,
+) -> Awaitable[OmadaGatewayPortStatus]:
+    return client.set_gateway_wan_port_connect_state(
+        port.port_number, enable, device, ipv6=ipv6
+    )
+
+
+GATEWAY_PORT_SWITCHES: list[GatewayPortSwitchEntityDescription] = [
+    GatewayPortSwitchEntityDescription(
+        key="wan_connect_ipv4",
+        translation_key="wan_connect_ipv4",
+        exists_func=lambda p: p.mode == GatewayPortMode.WAN,
+        set_func=partial(_wan_connect_disconnect, ipv6=False),
+        update_func=lambda p: p.wan_connected,
+    ),
+    GatewayPortSwitchEntityDescription(
+        key="wan_connect_ipv6",
+        translation_key="wan_connect_ipv6",
+        exists_func=lambda p: p.mode == GatewayPortMode.WAN and p.wan_ipv6_enabled,
+        set_func=partial(_wan_connect_disconnect, ipv6=True),
+        update_func=lambda p: p.ipv6_wan_connected,
+    ),
+]
+
+
 def get_port_base_name(port: OmadaSwitchPortDetails) -> str:
     """Get display name for a switch port."""
 
     if port.name == f"Port{port.port}":
-        return f"Port {port.port}"
+        return f"{port.port}"
-    return f"Port {port.port} ({port.name})"
+    return f"{port.port} ({port.name})"
 
 
 class OmadaNetworkSwitchPortPoEControl(
@@ -76,8 +143,9 @@ class OmadaNetworkSwitchPortPoEControl(
         self.port_details = coordinator.data[port_id]
         self.omada_client = coordinator.omada_client
         self._attr_unique_id = f"{device.mac}_{port_id}_poe"
-
+        self._attr_translation_placeholders = {
-        self._attr_name = f"{get_port_base_name(self.port_details)} PoE"
+            "port_name": get_port_base_name(self.port_details)
+        }
 
     async def async_added_to_hass(self) -> None:
         """When entity is added to hass."""
@@ -109,3 +177,72 @@ class OmadaNetworkSwitchPortPoEControl(
         """Handle updated data from the coordinator."""
         self.port_details = self.coordinator.data[self.port_id]
         self._refresh_state()
+
+
+class OmadaGatewayPortSwitchEntity(OmadaDeviceEntity[OmadaGateway], SwitchEntity):
+    """Generic toggle switch on a Gateway entity."""
+
+    _attr_has_entity_name = True
+    _port_details: OmadaGatewayPortStatus | None = None
+    entity_description: GatewayPortSwitchEntityDescription
+
+    def __init__(
+        self,
+        coordinator: OmadaGatewayCoordinator,
+        device: OmadaGateway,
+        port_number: int,
+        entity_description: GatewayPortSwitchEntityDescription,
+    ) -> None:
+        """Initialize the toggle switch."""
+        super().__init__(coordinator, device)
+        self.entity_description = entity_description
+        self._port_number = port_number
+        self._attr_unique_id = f"{device.mac}_{port_number}_{entity_description.key}"
+        self._attr_translation_placeholders = {"port_name": f"{port_number}"}
+
+    async def async_added_to_hass(self) -> None:
+        """When entity is added to hass."""
+        await super().async_added_to_hass()
+        self._do_update()
+
+    async def _async_turn_on_off(self, enable: bool) -> None:
+        if self._port_details:
+            self._port_details = await self.entity_description.set_func(
+                self.coordinator.omada_client, self.device, self._port_details, enable
+            )
+        self._attr_is_on = enable
+        # Refresh to make sure the requested changes stuck
+        await self.coordinator.async_request_refresh()
+
+    async def async_turn_on(self, **kwargs: Any) -> None:
+        """Turn the entity on."""
+        await self._async_turn_on_off(True)
+
+    async def async_turn_off(self, **kwargs: Any) -> None:
+        """Turn the entity off."""
+        await self._async_turn_on_off(False)
+
+    @property
+    def available(self) -> bool:
+        """Return true if entity is available."""
+        return bool(
+            super().available
+            and self._port_details
+            and self.entity_description.exists_func(self._port_details)
+        )
+
+    def _do_update(self) -> None:
+        gateway = self.coordinator.data[self.device.mac]
+
+        port = next(
+            p for p in gateway.port_status if p.port_number == self._port_number
+        )
+        if port:
+            self._port_details = port
+            self._attr_is_on = self.entity_description.update_func(port)
+
+    @callback
+    def _handle_coordinator_update(self) -> None:
+        """Handle updated data from the coordinator."""
+        self._do_update()
+        self.async_write_ha_state()

tests/components/tplink_omada/conftest.py

@@ -5,7 +5,12 @@ import json
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
-from tplink_omada_client.devices import OmadaSwitch, OmadaSwitchPortDetails
+from tplink_omada_client.devices import (
+    OmadaGateway,
+    OmadaListDevice,
+    OmadaSwitch,
+    OmadaSwitchPortDetails,
+)
 
 from homeassistant.components.tplink_omada.config_flow import CONF_SITE
 from homeassistant.components.tplink_omada.const import DOMAIN
@@ -45,10 +50,19 @@ def mock_setup_entry() -> Generator[AsyncMock, None, None]:
 def mock_omada_site_client() -> Generator[AsyncMock, None, None]:
     """Mock Omada site client."""
     site_client = AsyncMock()
+
+    gateway_data = json.loads(load_fixture("gateway-TL-ER7212PC.json", DOMAIN))
+    gateway = OmadaGateway(gateway_data)
+    site_client.get_gateway.return_value = gateway
+
     switch1_data = json.loads(load_fixture("switch-TL-SG3210XHP-M2.json", DOMAIN))
     switch1 = OmadaSwitch(switch1_data)
     site_client.get_switches.return_value = [switch1]
 
+    devices_data = json.loads(load_fixture("devices.json", DOMAIN))
+    devices = [OmadaListDevice(d) for d in devices_data]
+    site_client.get_devices.return_value = devices
+
     switch1_ports_data = json.loads(
         load_fixture("switch-ports-TL-SG3210XHP-M2.json", DOMAIN)
     )

tests/components/tplink_omada/fixtures/devices.json

@@ -0,0 +1,41 @@
+[
+  {
+    "type": "gateway",
+    "mac": "AA-BB-CC-DD-EE-FF",
+    "name": "Test Router",
+    "model": "ER7212PC",
+    "compoundModel": "ER7212PC v1.0",
+    "showModel": "ER7212PC v1.0",
+    "firmwareVersion": "1.1.1 Build 20230901 Rel.55651",
+    "version": "1.1.1",
+    "hwVersion": "ER7212PC v1.0",
+    "uptime": "32day(s) 5h 39m 27s",
+    "uptimeLong": 2785167,
+    "cpuUtil": 16,
+    "memUtil": 47,
+    "status": 14,
+    "statusCategory": 1,
+    "site": "Test",
+    "needUpgrade": false
+  },
+  {
+    "type": "switch",
+    "mac": "54-AF-97-00-00-01",
+    "name": "Test PoE Switch",
+    "model": "TL-SG3210XHP-M2",
+    "modelVersion": "1.0",
+    "compoundModel": "TL-SG3210XHP-M2 v1.0",
+    "showModel": "TL-SG3210XHP-M2 v1.0",
+    "firmwareVersion": "1.0.12 Build 20230203 Rel.36545",
+    "version": "1.0.12",
+    "hwVersion": "1.0",
+    "uptime": "1day(s) 8h 27m 26s",
+    "uptimeLong": 116846,
+    "cpuUtil": 10,
+    "memUtil": 20,
+    "status": 14,
+    "statusCategory": 1,
+    "needUpgrade": false,
+    "fwDownload": false
+  }
+]

tests/components/tplink_omada/fixtures/switch-TL-SG3210XHP-M2.json

@@ -11,7 +11,7 @@
   "hwVersion": "1.0",
   "status": 14,
   "statusCategory": 1,
-  "site": "000000000000000000000000",
+  "site": "Test",
   "omadacId": "00000000000000000000000000000000",
   "compatible": 0,
   "sn": "Y220000000001",
@@ -124,7 +124,7 @@
     {
       "id": "000000000000000000000002",
       "port": 2,
-      "name": "Port2",
+      "name": "Renamed Port",
       "disable": false,
       "type": 1,
       "maxSpeed": 4,

tests/components/tplink_omada/fixtures/switch-ports-TL-SG3210XHP-M2.json

@@ -108,7 +108,7 @@
     "switchId": "640934810000000000000000",
     "switchMac": "54-AF-97-00-00-01",
     "site": "000000000000000000000000",
-    "name": "Port2",
+    "name": "Renamed Port",
     "disable": false,
     "type": 1,
     "maxSpeed": 4,

tests/components/tplink_omada/snapshots/test_switch.ambr

@@ -1,4 +1,53 @@
 # serializer version: 1
+# name: test_gateway_api_fail_disables_switch_entities
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'friendly_name': 'Test Router Port 4 Internet Connected',
+    }),
+    'context': <ANY>,
+    'entity_id': 'switch.test_router_port_4_internet_connected',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': 'on',
+  })
+# ---
+# name: test_gateway_connect_ipv4_switch
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'friendly_name': 'Test Router Port 4 Internet Connected',
+    }),
+    'context': <ANY>,
+    'entity_id': 'switch.test_router_port_4_internet_connected',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': 'on',
+  })
+# ---
+# name: test_gateway_disappear_disables_switches
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'friendly_name': 'Test Router Port 4 Internet Connected',
+      'icon': 'mdi:ethernet',
+    }),
+    'context': <ANY>,
+    'entity_id': 'switch.test_router_port_4_internet_connected',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': 'on',
+  })
+# ---
+# name: test_gateway_port_change_disables_switch_entities
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'friendly_name': 'Test Router Port 4 Internet Connected',
+    }),
+    'context': <ANY>,
+    'entity_id': 'switch.test_router_port_4_internet_connected',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': 'on',
+  })
+# ---
 # name: test_poe_switches
   StateSnapshot({
     'attributes': ReadOnlyDict({
@@ -182,10 +231,10 @@
 # name: test_poe_switches.2
   StateSnapshot({
     'attributes': ReadOnlyDict({
-      'friendly_name': 'Test PoE Switch Port 2 PoE',
+      'friendly_name': 'Test PoE Switch Port 2 (Renamed Port) PoE',
     }),
     'context': <ANY>,
-    'entity_id': 'switch.test_poe_switch_port_2_poe',
+    'entity_id': 'switch.test_poe_switch_port_2_renamed_port_poe',
     'last_changed': <ANY>,
     'last_updated': <ANY>,
     'state': 'on',
@@ -203,7 +252,7 @@
     'disabled_by': None,
     'domain': 'switch',
     'entity_category': <EntityCategory.CONFIG: 'config'>,
-    'entity_id': 'switch.test_poe_switch_port_2_poe',
+    'entity_id': 'switch.test_poe_switch_port_2_renamed_port_poe',
     'has_entity_name': True,
     'hidden_by': None,
     'icon': None,
@@ -215,7 +264,7 @@
     }),
     'original_device_class': None,
     'original_icon': None,
-    'original_name': 'Port 2 PoE',
+    'original_name': 'Port 2 (Renamed Port) PoE',
     'platform': 'tplink_omada',
     'previous_unique_id': None,
     'supported_features': 0,

tests/components/tplink_omada/test_switch.py

@@ -1,18 +1,30 @@
 """Tests for TP-Link Omada switch entities."""
-
+from datetime import timedelta
+from typing import Any
 from unittest.mock import MagicMock
 
 from syrupy.assertion import SnapshotAssertion
 from tplink_omada_client import SwitchPortOverrides
 from tplink_omada_client.definitions import PoEMode
-from tplink_omada_client.devices import OmadaSwitch, OmadaSwitchPortDetails
+from tplink_omada_client.devices import (
+    OmadaGateway,
+    OmadaGatewayPortStatus,
+    OmadaSwitch,
+    OmadaSwitchPortDetails,
+)
+from tplink_omada_client.exceptions import InvalidDevice
 
 from homeassistant.components import switch
+from homeassistant.components.tplink_omada.controller import POLL_GATEWAY
 from homeassistant.const import ATTR_ENTITY_ID
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry as er
+from homeassistant.util.dt import utcnow
 
-from tests.common import MockConfigEntry
+from tests.common import MockConfigEntry, async_fire_time_changed
+
+UPDATE_INTERVAL = timedelta(seconds=10)
+POLL_INTERVAL = timedelta(seconds=POLL_GATEWAY + 10)
 
 
 async def test_poe_switches(
@@ -23,15 +35,124 @@ async def test_poe_switches(
 ) -> None:
     """Test PoE switch."""
     poe_switch_mac = "54-AF-97-00-00-01"
-    for i in range(1, 9):
+    await _test_poe_switch(
-        await _test_poe_switch(
+        hass,
-            hass,
+        mock_omada_site_client,
-            mock_omada_site_client,
+        "switch.test_poe_switch_port_1_poe",
-            f"switch.test_poe_switch_port_{i}_poe",
+        poe_switch_mac,
-            poe_switch_mac,
+        1,
-            i,
+        snapshot,
-            snapshot,
+    )
+
+    await _test_poe_switch(
+        hass,
+        mock_omada_site_client,
+        "switch.test_poe_switch_port_2_renamed_port_poe",
+        poe_switch_mac,
+        2,
+        snapshot,
+    )
+
+
+async def test_gateway_connect_ipv4_switch(
+    hass: HomeAssistant,
+    mock_omada_site_client: MagicMock,
+    init_integration: MockConfigEntry,
+    snapshot: SnapshotAssertion,
+) -> None:
+    """Test gateway connected switches."""
+    gateway_mac = "AA-BB-CC-DD-EE-FF"
+
+    entity_id = "switch.test_router_port_4_internet_connected"
+    entity = hass.states.get(entity_id)
+    assert entity == snapshot
+
+    test_gateway = await mock_omada_site_client.get_gateway(gateway_mac)
+    port_status = test_gateway.port_status[3]
+    assert port_status.port_number == 4
+
+    mock_omada_site_client.set_gateway_wan_port_connect_state.reset_mock()
+    mock_omada_site_client.set_gateway_wan_port_connect_state.return_value = (
+        _get_updated_gateway_port_status(
+            mock_omada_site_client, test_gateway, 3, "internetState", 0
         )
+    )
+    await call_service(hass, "turn_off", entity_id)
+    mock_omada_site_client.set_gateway_wan_port_connect_state.assert_called_once_with(
+        4, False, test_gateway, ipv6=False
+    )
+
+    async_fire_time_changed(hass, utcnow() + UPDATE_INTERVAL)
+    await hass.async_block_till_done()
+
+    entity = hass.states.get(entity_id)
+    assert entity.state == "off"
+
+    mock_omada_site_client.set_gateway_wan_port_connect_state.reset_mock()
+    mock_omada_site_client.set_gateway_wan_port_connect_state.return_value = (
+        _get_updated_gateway_port_status(
+            mock_omada_site_client, test_gateway, 3, "internetState", 1
+        )
+    )
+    await call_service(hass, "turn_on", entity_id)
+    mock_omada_site_client.set_gateway_wan_port_connect_state.assert_called_once_with(
+        4, True, test_gateway, ipv6=False
+    )
+
+    async_fire_time_changed(hass, utcnow() + UPDATE_INTERVAL)
+    await hass.async_block_till_done()
+
+    entity = hass.states.get(entity_id)
+    assert entity.state == "on"
+
+
+async def test_gateway_api_fail_disables_switch_entities(
+    hass: HomeAssistant,
+    mock_omada_site_client: MagicMock,
+    init_integration: MockConfigEntry,
+    snapshot: SnapshotAssertion,
+) -> None:
+    """Test gateway connected switches."""
+    entity_id = "switch.test_router_port_4_internet_connected"
+    entity = hass.states.get(entity_id)
+    assert entity == snapshot
+    assert entity.state == "on"
+
+    mock_omada_site_client.get_gateway.reset_mock()
+    mock_omada_site_client.get_gateway.side_effect = InvalidDevice("Expected error")
+
+    async_fire_time_changed(hass, utcnow() + POLL_INTERVAL)
+    await hass.async_block_till_done()
+
+    entity = hass.states.get(entity_id)
+    assert entity.state == "unavailable"
+
+
+async def test_gateway_port_change_disables_switch_entities(
+    hass: HomeAssistant,
+    mock_omada_site_client: MagicMock,
+    init_integration: MockConfigEntry,
+    snapshot: SnapshotAssertion,
+) -> None:
+    """Test gateway connected switch reconfigure."""
+
+    gateway_mac = "AA-BB-CC-DD-EE-FF"
+    test_gateway = await mock_omada_site_client.get_gateway(gateway_mac)
+
+    entity_id = "switch.test_router_port_4_internet_connected"
+    entity = hass.states.get(entity_id)
+    assert entity == snapshot
+    assert entity.state == "on"
+
+    mock_omada_site_client.get_gateway.reset_mock()
+    # Set Port 4 to LAN mode
+    _get_updated_gateway_port_status(mock_omada_site_client, test_gateway, 3, "mode", 1)
+
+    async_fire_time_changed(hass, utcnow() + POLL_INTERVAL)
+    await hass.async_block_till_done()
+
+    entity = hass.states.get(entity_id)
+    assert entity.state == "unavailable"
 
 
 async def _test_poe_switch(
@@ -94,6 +215,7 @@ async def _test_poe_switch(
         True,
         **mock_omada_site_client.update_switch_port.call_args.kwargs,
     )
+    await hass.async_block_till_done()
     entity = hass.states.get(entity_id)
     assert entity.state == "on"
 
@@ -116,6 +238,20 @@ async def _update_port_details(
     return OmadaSwitchPortDetails(raw_data)
 
 
+def _get_updated_gateway_port_status(
+    mock_omada_site_client: MagicMock,
+    gateway: OmadaGateway,
+    port: int,
+    key: str,
+    value: Any,
+) -> OmadaGatewayPortStatus:
+    gateway_data = gateway.raw_data.copy()
+    gateway_data["portStats"][port][key] = value
+    mock_omada_site_client.get_gateway.reset_mock()
+    mock_omada_site_client.get_gateway.return_value = OmadaGateway(gateway_data)
+    return OmadaGatewayPortStatus(gateway_data["portStats"][port])
+
+
 def call_service(hass: HomeAssistant, service: str, entity_id: str):
     """Call any service on entity."""
     return hass.services.async_call(