From ef4515ed70e8d92ae471e348a5d094d3879911fc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ville=20Skytt=C3=A4?= <ville.skytta@iki.fi>
Date: Sun, 8 Dec 2019 13:21:48 +0200
Subject: [PATCH] Add Huawei LTE reboot and clear traffic statistics services
 (#29594)

* Add clear traffic statistics service

* Add reboot service

* Register services as admin ones

* Make URL optional when there's only one router configured

* Eliminate one if/else indent level

* Remove unnecessary .keys() with sorted()
---
 .../components/huawei_lte/__init__.py         | 37 +++++++++++++++++++
 homeassistant/components/huawei_lte/const.py  |  3 ++
 .../components/huawei_lte/services.yaml       | 13 +++++++
 3 files changed, 53 insertions(+)
 create mode 100644 homeassistant/components/huawei_lte/services.yaml

diff --git a/homeassistant/components/huawei_lte/__init__.py b/homeassistant/components/huawei_lte/__init__.py
index ada1f0a8abd..30b8eb1bc6d 100644
--- a/homeassistant/components/huawei_lte/__init__.py
+++ b/homeassistant/components/huawei_lte/__init__.py
@@ -58,6 +58,8 @@ from .const import (
     KEY_MONITORING_STATUS,
     KEY_MONITORING_TRAFFIC_STATISTICS,
     KEY_WLAN_HOST_LIST,
+    SERVICE_CLEAR_TRAFFIC_STATISTICS,
+    SERVICE_REBOOT,
     UPDATE_OPTIONS_SIGNAL,
     UPDATE_SIGNAL,
 )
@@ -103,6 +105,8 @@ CONFIG_SCHEMA = vol.Schema(
     extra=vol.ALLOW_EXTRA,
 )
 
+SERVICE_SCHEMA = vol.Schema({vol.Optional(CONF_URL): cv.url})
+
 CONFIG_ENTRY_PLATFORMS = (
     BINARY_SENSOR_DOMAIN,
     DEVICE_TRACKER_DOMAIN,
@@ -387,6 +391,39 @@ async def async_setup(hass: HomeAssistantType, config) -> bool:
     for router_config in config.get(DOMAIN, []):
         domain_config[url_normalize(router_config.pop(CONF_URL))] = router_config
 
+    def service_handler(service) -> None:
+        """Apply a service."""
+        url = service.data.get(CONF_URL)
+        routers = hass.data[DOMAIN].routers
+        if url:
+            router = routers.get(url)
+        elif len(routers) == 1:
+            router = next(iter(routers.values()))
+        else:
+            _LOGGER.error(
+                "%s: more than one router configured, must specify one of URLs %s",
+                service.service,
+                sorted(routers),
+            )
+            return
+        if not router:
+            _LOGGER.error("%s: router %s unavailable", service.service, url)
+            return
+
+        if service.service == SERVICE_CLEAR_TRAFFIC_STATISTICS:
+            result = router.client.monitoring.set_clear_traffic()
+            _LOGGER.debug("%s: %s", service.service, result)
+        elif service.service == SERVICE_REBOOT:
+            result = router.client.device.reboot()
+            _LOGGER.debug("%s: %s", service.service, result)
+        else:
+            _LOGGER.error("%s: unsupported service", service.service)
+
+    for service in (SERVICE_CLEAR_TRAFFIC_STATISTICS, SERVICE_REBOOT):
+        hass.helpers.service.async_register_admin_service(
+            DOMAIN, service, service_handler, schema=SERVICE_SCHEMA,
+        )
+
     for url, router_config in domain_config.items():
         hass.async_create_task(
             hass.config_entries.flow.async_init(
diff --git a/homeassistant/components/huawei_lte/const.py b/homeassistant/components/huawei_lte/const.py
index b6e079576ac..c71b51435e1 100644
--- a/homeassistant/components/huawei_lte/const.py
+++ b/homeassistant/components/huawei_lte/const.py
@@ -12,6 +12,9 @@ UNIT_SECONDS = "s"
 
 CONNECTION_TIMEOUT = 10
 
+SERVICE_CLEAR_TRAFFIC_STATISTICS = "clear_traffic_statistics"
+SERVICE_REBOOT = "reboot"
+
 KEY_DEVICE_BASIC_INFORMATION = "device_basic_information"
 KEY_DEVICE_INFORMATION = "device_information"
 KEY_DEVICE_SIGNAL = "device_signal"
diff --git a/homeassistant/components/huawei_lte/services.yaml b/homeassistant/components/huawei_lte/services.yaml
new file mode 100644
index 00000000000..428745ee33e
--- /dev/null
+++ b/homeassistant/components/huawei_lte/services.yaml
@@ -0,0 +1,13 @@
+clear_traffic_statistics:
+  description: Clear traffic statistics.
+  fields:
+    url:
+      description: URL of router to clear; optional when only one is configured.
+      example: http://192.168.100.1/
+
+reboot:
+  description: Reboot router.
+  fields:
+    url:
+      description: URL of router to reboot; optional when only one is configured.
+      example: http://192.168.100.1/