"""/speed: an Ookla Speedtest run from the bot's own network. The official Speedtest CLI is used when it is installed. If it is missing or fails, the ``speedtest-cli`` Python package (also an Ookla client) is tried. Both give the same figures; only the parsing differs. """ from __future__ import annotations import asyncio import html import json import logging import os import shutil from dataclasses import dataclass from typing import Any, Dict, List, Optional from .utils import esc log = logging.getLogger(__name__) CLI_TIMEOUT = 120 # seconds for one run of either client CANDIDATES = ("ookla-speedtest", "speedtest") # the Dockerfile installs the official binary as ookla-speedtest class SpeedError(RuntimeError): """The speed test could not be run; the message is safe to show to the user.""" @dataclass class SpeedResult: download_mbps: float upload_mbps: float ping_ms: float jitter_ms: Optional[float] = None packet_loss: Optional[float] = None isp: str = "" server: str = "" location: str = "" url: str = "" # the speedtest.net page of this result source: str = "" def _num(value: Any) -> Optional[float]: try: return float(value) except (TypeError, ValueError): return None def parse_cli(data: Dict[str, Any]) -> SpeedResult: """The JSON the official CLI prints (bandwidth is in bytes per second).""" download = _num((data.get("download") or {}).get("bandwidth")) upload = _num((data.get("upload") or {}).get("bandwidth")) if download is None or upload is None: raise SpeedError("the Speedtest CLI gave no download or upload figure") ping = data.get("ping") or {} server = data.get("server") or {} location = ", ".join(part for part in (server.get("location"), server.get("country")) if part) return SpeedResult( download_mbps=download * 8 / 1e6, upload_mbps=upload * 8 / 1e6, ping_ms=_num(ping.get("latency")) or 0.0, jitter_ms=_num(ping.get("jitter")), packet_loss=_num(data.get("packetLoss")), isp=str(data.get("isp") or ""), server=str(server.get("name") or server.get("host") or ""), location=location, url=str((data.get("result") or {}).get("url") or ""), source="Ookla Speedtest CLI", ) def parse_python(data: Dict[str, Any]) -> SpeedResult: """``Speedtest.results.dict()`` of the speedtest-cli package (speeds are in bits per second).""" download, upload = _num(data.get("download")), _num(data.get("upload")) if download is None or upload is None: raise SpeedError("speedtest-cli gave no download or upload figure") server = data.get("server") or {} client = data.get("client") or {} location = ", ".join(part for part in (server.get("name"), server.get("country")) if part) return SpeedResult( download_mbps=download / 1e6, upload_mbps=upload / 1e6, ping_ms=_num(data.get("ping")) or 0.0, isp=str(client.get("isp") or ""), server=str(server.get("sponsor") or ""), location=location, url=str(data.get("share") or ""), source="speedtest-cli (Ookla)", ) def _last_result(stdout: str) -> Optional[Dict[str, Any]]: """The ``{"type": "result", ...}`` line of the CLI's JSON output.""" found = None for line in stdout.splitlines(): line = line.strip() if not line.startswith("{"): continue try: obj = json.loads(line) except ValueError: continue if isinstance(obj, dict) and obj.get("type") == "result": found = obj return found def _log_message(stdout: str, stderr: str) -> str: """Why the CLI failed, as far as it says.""" for text in (stdout, stderr): for line in reversed(text.splitlines()): line = line.strip() if line.startswith("{"): try: obj = json.loads(line) except ValueError: continue if isinstance(obj, dict) and obj.get("message"): return str(obj["message"]) elif line: return line return "no output" async def _run(cmd: List[str], timeout: float) -> "tuple[int, str, str]": try: proc = await asyncio.create_subprocess_exec( *cmd, stdin=asyncio.subprocess.DEVNULL, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE ) except (FileNotFoundError, PermissionError, OSError) as exc: raise SpeedError(f"{cmd[0]} cannot be started ({exc})") from exc try: out, err = await asyncio.wait_for(proc.communicate(), timeout) except asyncio.TimeoutError: proc.kill() await proc.wait() raise SpeedError("the speed test took too long") from None return proc.returncode or 0, out.decode("utf-8", "replace"), err.decode("utf-8", "replace") async def find_ookla_cli() -> Optional[str]: """The official CLI, if installed. ``speedtest`` can also be the Python package's script, so ``--version`` decides.""" names = [os.environ.get("SPEEDTEST_BIN", "")] + list(CANDIDATES) for name in dict.fromkeys(n for n in names if n): path = shutil.which(name) if not path: continue try: code, out, _err = await _run([path, "--version"], 15) except SpeedError: continue if code == 0 and "ookla" in out.lower(): return path return None async def run_cli(path: str, timeout: float = CLI_TIMEOUT) -> SpeedResult: code, out, err = await _run([path, "--accept-license", "--accept-gdpr", "--format=json"], timeout) result = _last_result(out) if result is None: raise SpeedError(f"the Speedtest CLI failed: {_log_message(out, err)}") return parse_cli(result) def _run_python() -> Dict[str, Any]: """Blocking: one test with the speedtest-cli package.""" try: import speedtest # type: ignore except ImportError as exc: raise SpeedError("the speedtest-cli package is not installed") from exc try: test = speedtest.Speedtest(secure=True) test.get_best_server() test.download() test.upload() return dict(test.results.dict()) except Exception as exc: # noqa: BLE001 - the package raises its own family of errors raise SpeedError(f"speedtest-cli failed: {exc}") from exc async def run_speedtest() -> SpeedResult: """Run a test with the official CLI, or with the Python client when the CLI is missing or fails.""" problems: List[str] = [] cli = await find_ookla_cli() if cli is not None: try: return await run_cli(cli) except SpeedError as exc: log.warning("Speedtest CLI failed: %s", exc) problems.append(str(exc)) else: problems.append("the Speedtest CLI is not installed") try: data = await asyncio.wait_for(asyncio.to_thread(_run_python), CLI_TIMEOUT) return parse_python(data) except asyncio.TimeoutError: problems.append("speedtest-cli took too long") except SpeedError as exc: problems.append(str(exc)) raise SpeedError("; ".join(problems)) def _mbps(value: float) -> str: return f"{value:,.2f} Mbps" def format_result(r: SpeedResult) -> str: lines = [ "⚡ SPEEDTEST · Ookla", "", f"◈ Download ▸ {_mbps(r.download_mbps)}", f"◈ Upload ▸ {_mbps(r.upload_mbps)}", f"◈ Ping ▸ {r.ping_ms:.2f} ms" + (f" · jitter {r.jitter_ms:.2f} ms" if r.jitter_ms is not None else ""), ] if r.packet_loss is not None: lines.append(f"◈ Packet loss ▸ {r.packet_loss:g}%") if r.isp: lines.append(f"◈ ISP ▸ {esc(r.isp)}") server = " · ".join(part for part in (esc(r.server), esc(r.location)) if part) if server: lines.append(f"◈ Server ▸ {server}") if r.url.startswith(("http://", "https://")): lines.append(f"◈ Result ▸ speedtest.net") lines.append(f"\nMeasured from the machine the bot runs on · {esc(r.source)}") return "\n".join(lines)