Add speedtest data collection (#393)

asusrouter/asusrouter.py

@@ -243,6 +243,13 @@ async def async_get_identity(self, force: bool = False) -> AsusDevice:
                 remove_data_rule(AsusData.WIREGUARD_CLIENT)
                 remove_data_rule(AsusData.WIREGUARD_SERVER)
 
+            # Ookla Speedtest
+            if self._identity.ookla is False:
+                remove_data_rule(AsusData.SPEEDTEST)
+                # remove_data_rule(AsusData.SPEEDTEST_HISTORY)
+                remove_data_rule(AsusData.SPEEDTEST_RESULT)
+                # remove_data_rule(AsusData.SPEEDTEST_SERVERS)
+
         # Return new identity
         return self._identity
 

asusrouter/modules/data.py

@@ -30,6 +30,10 @@ class AsusData(str, Enum):
     PORT_FORWARDING = "port_forwarding"
     PORTS = "ports"
     RAM = "ram"
+    SPEEDTEST = "speedtest"
+    # SPEEDTEST_HISTORY = "speedtest_history"
+    SPEEDTEST_RESULT = "speedtest_result"
+    # SPEEDTEST_SERVERS = "speedtest_servers"
     SYSINFO = "sysinfo"
     SYSTEM = "system"
     TEMPERATURE = "temperature"

asusrouter/modules/data_finder.py

@@ -82,6 +82,15 @@ def __init__(
         ("memory_usage", "appobj"),
         ("netdev", "appobj"),
     ],
+    "speedtest": [
+        ("ookla_speedtest_get_result", ""),
+    ],
+    # "speedtest_history": [
+    #     ("ookla_speedtest_get_history", ""),
+    # ],
+    # "speedtest_servers": [
+    #     ("ookla_speedtest_get_servers", ""),
+    # ],
     "vpnc": [
         ("get_vpnc_status", ""),
     ],
@@ -111,6 +120,7 @@ def __init__(
         "vts_rulelist",
         "vts_enable_x",
     ],
+    "speedtest": ["ookla_state"],
     "vpnc": [
         "vpnc_clientlist",
     ],
@@ -199,6 +209,18 @@ def __init__(
     ),
     AsusData.PORTS: AsusDataFinder([Endpoint.PORT_STATUS, Endpoint.ETHERNET_PORTS]),
     AsusData.RAM: AsusData.CPU,
+    AsusData.SPEEDTEST: AsusDataFinder(
+        Endpoint.HOOK,
+        nvram=ASUSDATA_NVRAM["speedtest"],
+        request=ASUSDATA_REQUEST["speedtest"],
+    ),
+    # AsusData.SPEEDTEST_HISTORY: AsusDataFinder(
+    #     Endpoint.HOOK, request=ASUSDATA_REQUEST["speedtest_history"]
+    # ),
+    AsusData.SPEEDTEST_RESULT: AsusData.SPEEDTEST,
+    # AsusData.SPEEDTEST_SERVERS: AsusDataFinder(
+    #     Endpoint.HOOK, request=ASUSDATA_REQUEST["speedtest_servers"]
+    # ),
     AsusData.SYSINFO: AsusDataFinder(Endpoint.SYSINFO),
     AsusData.TEMPERATURE: AsusDataFinder(Endpoint.TEMPERATURE),
     AsusData.VPNC: AsusDataFinder(

asusrouter/modules/endpoint/hook.py

@@ -31,6 +31,7 @@
 from asusrouter.tools.converters import (
     run_method,
     safe_bool,
+    safe_datetime,
     safe_int,
     safe_return,
     safe_speed,
@@ -44,6 +45,7 @@
 from .hook_const import (
     MAP_NETWORK,
     MAP_OVPN_SERVER_388,
+    MAP_SPEEDTEST,
     MAP_VPNC_WIREGUARD,
     MAP_WAN,
     MAP_WAN_ITEM,
@@ -133,6 +135,12 @@ def process(data: dict[str, Any]) -> dict[AsusData, Any]:
         memory_usage = data.get("memory_usage", {})
         state[AsusData.RAM] = process_ram(memory_usage) if memory_usage else {}
 
+    # Speedtest
+    if "ookla_state" in data:
+        speedtest = process_speedtest(data)
+        state[AsusData.SPEEDTEST_RESULT] = speedtest.get("result")
+        state[AsusData.SPEEDTEST] = speedtest.get("data")
+
     # VPNC
     if "vpnc_clientlist" in data:
         vpnc, vpnc_clientlist = process_vpnc(data)
@@ -406,6 +414,64 @@ def process_ram(memory_usage: dict[str, Any]) -> dict[str, Any]:
     return ram
 
 
+def process_speedtest(data: dict[str, Any]) -> dict[str, Any]:
+    """Process Speedtest data."""
+
+    speedtest: dict[str, Any] = {}
+
+    # Convert the data
+    for keys in MAP_SPEEDTEST:
+        key, key_to_use, method = safe_unpack_keys(keys)
+        state_value = data.get(key)
+        if state_value:
+            speedtest[key_to_use] = run_method(state_value, method)
+
+    # Get detailed result
+    test_result = speedtest.pop("result", None)
+
+    # Get the last speedtest step
+    last_run_step: Optional[dict[str, Any]] = process_speedtest_last_step(test_result)
+
+    # Save the last tested data directly in the speedtest dict
+    for key, value in last_run_step.items():
+        speedtest[key] = value
+
+    # Convert the timestamp to UTC
+    speedtest["timestamp"] = safe_datetime(speedtest.get("timestamp"))
+
+    # Rename servers list
+    if "download" in speedtest:
+        speedtest["download"]["server_list"] = speedtest["download"].pop("servers", None)
+
+    # Remove extra values
+    speedtest.pop("type", None)
+
+    return {
+        "data": speedtest,
+        "result": test_result,
+    }
+
+
+def process_speedtest_last_step(data: Optional[list[dict[str, Any]]]) -> dict[str, Any]:
+    """Process Speedtest progress data."""
+
+    last_step: dict[str, Any] = {}
+
+    # Check the length of the data
+    if not data or len(data) < 2:
+        return last_step
+
+    # Check the last step
+    last_step = data[-2]
+
+    # If the last step is a result, then the test is finished
+    # and we already have all the data
+    if last_step.get("type") != "result":
+        return {}
+
+    return last_step
+
+
 def process_vpnc(data: dict[str, Any]) -> Tuple[dict[AsusVPNType, dict[int, Any]], str]:
     """Process VPNC data."""
 

asusrouter/modules/endpoint/hook_const.py

@@ -12,6 +12,7 @@
     safe_int,
     safe_list_csv,
     safe_list_from_string,
+    safe_timestamp_to_utc,
 )
 
 MAP_NETWORK: dict[str, str] = {
@@ -59,6 +60,13 @@
     ("VPNServer_enable", "state", [safe_int, AsusOVPNServer]),
 )
 
+MAP_SPEEDTEST = (
+    ("ookla_state", "state", safe_int),
+    ("ookla_speedtest_get_history", "history"),
+    ("ookla_speedtest_get_servers", "servers"),
+    ("ookla_speedtest_get_result", "result"),
+)
+
 MAP_WAN = (
     ("get_wan_unit", "unit", safe_int),
     ("link_internet", "link", [safe_int, ConnectionStatus]),

asusrouter/modules/identity.py

@@ -58,9 +58,10 @@ class AsusDevice:  # pylint: disable=too-many-instance-attributes
     services: Optional[list[str]] = None
 
     # Flags for device features
+    aura: bool = False
     led: bool = False
     ledg: bool = False
-    aura: bool = False
+    ookla: bool = False
     vpn_status: bool = False
 
 
@@ -159,6 +160,10 @@ def _read_nvram(data: dict[str, Any]) -> dict[str, Any]:
         if value in WLAN_TYPE:
             identity["wlan"].append(WLAN_TYPE[value])
 
+    # OOKLA Speedtest
+    if "ookla" in identity["services"]:
+        identity["ookla"] = True
+
     return identity