Create speedtest.py

bot/core/speedtest.py

@@ -0,0 +1,170 @@
+# PATH: bot/core/speedtest.py
+import asyncio
+import time
+import shutil
+from typing import Dict, Any, Optional
+
+import httpx
+
+from bot.core.uptime import uptime_seconds
+from bot.core.progress import human_bytes, human_eta
+
+DL_URL = "https://speed.cloudflare.com/__down?bytes={bytes}"
+UL_URL = "https://httpbin.org/post"  # simple echo endpoint (kabhi kabhi rate-limit ho sakta)
+
+def _bps_to_mbps(bps: float) -> float:
+    return (bps * 8.0) / 1_000_000.0
+
+async def tcp_ping_ms(host: str = "1.1.1.1", port: int = 443, timeout: float = 3.0) -> Optional[float]:
+    t0 = time.perf_counter()
+    try:
+        conn = asyncio.open_connection(host, port)
+        r, w = await asyncio.wait_for(conn, timeout=timeout)
+        w.close()
+        try:
+            await w.wait_closed()
+        except Exception:
+            pass
+        return (time.perf_counter() - t0) * 1000.0
+    except Exception:
+        return None
+
+async def download_test(bytes_to_dl: int = 50 * 1024 * 1024, timeout: float = 30.0) -> Dict[str, Any]:
+    url = DL_URL.format(bytes=bytes_to_dl)
+    t0 = time.perf_counter()
+    total = 0
+
+    async with httpx.AsyncClient(follow_redirects=True, timeout=timeout) as client:
+        async with client.stream("GET", url) as resp:
+            resp.raise_for_status()
+            async for chunk in resp.aiter_bytes():
+                total += len(chunk)
+
+    dt = max(0.0001, time.perf_counter() - t0)
+    bps = total / dt
+    return {
+        "ok": True,
+        "bytes": total,
+        "sec": dt,
+        "bps": bps,
+        "mbps": _bps_to_mbps(bps),
+    }
+
+async def upload_test(bytes_to_ul: int = 10 * 1024 * 1024, timeout: float = 30.0) -> Dict[str, Any]:
+    # stream upload (no big RAM spike)
+    chunk = b"\x00" * (256 * 1024)
+    remaining = bytes_to_ul
+
+    async def gen():
+        nonlocal remaining
+        while remaining > 0:
+            take = chunk if remaining >= len(chunk) else b"\x00" * remaining
+            remaining -= len(take)
+            yield take
+
+    t0 = time.perf_counter()
+    async with httpx.AsyncClient(timeout=timeout) as client:
+        r = await client.post(
+            UL_URL,
+            content=gen(),
+            headers={"content-type": "application/octet-stream"},
+        )
+        r.raise_for_status()
+
+    sent = bytes_to_ul
+    dt = max(0.0001, time.perf_counter() - t0)
+    bps = sent / dt
+    return {
+        "ok": True,
+        "bytes": sent,
+        "sec": dt,
+        "bps": bps,
+        "mbps": _bps_to_mbps(bps),
+    }
+
+def storage_info(path: str = ".") -> Dict[str, Any]:
+    du = shutil.disk_usage(path)
+    return {
+        "total": du.total,
+        "used": du.used,
+        "free": du.free,
+    }
+
+async def run_speedtest(worker2_stats_fn=None) -> Dict[str, Any]:
+    out: Dict[str, Any] = {}
+    out["uptime_sec"] = uptime_seconds()
+
+    # ping (tcp)
+    out["ping_ms"] = await tcp_ping_ms()
+
+    # storage
+    st = storage_info(".")
+    out["storage_total"] = st["total"]
+    out["storage_free"] = st["free"]
+
+    # download
+    try:
+        out["download"] = await download_test()
+    except Exception as e:
+        out["download"] = {"ok": False, "err": f"{type(e).__name__}: {e}"}
+
+    # upload
+    try:
+        out["upload"] = await upload_test()
+    except Exception as e:
+        out["upload"] = {"ok": False, "err": f"{type(e).__name__}: {e}"}
+
+    # optional: Worker2 KV RTT (stats_today hits KV)
+    if worker2_stats_fn:
+        t0 = time.perf_counter()
+        try:
+            j = await worker2_stats_fn()
+            out["w2_kv_rtt_ms"] = (time.perf_counter() - t0) * 1000.0
+            out["w2_ok"] = bool(isinstance(j, dict) and j.get("ok"))
+        except Exception as e:
+            out["w2_kv_rtt_ms"] = None
+            out["w2_ok"] = False
+            out["w2_err"] = f"{type(e).__name__}: {e}"
+
+    return out
+
+def format_report(r: Dict[str, Any]) -> str:
+    up = r.get("uptime_sec", 0)
+    ping = r.get("ping_ms", None)
+
+    st_total = r.get("storage_total", 0)
+    st_free = r.get("storage_free", 0)
+
+    dl = r.get("download", {})
+    ul = r.get("upload", {})
+
+    lines = []
+    lines.append("🏁 **SpeedTest Report**")
+    lines.append("")
+    lines.append(f"🕒 Uptime: `{human_eta(up)}`")
+    lines.append(f"💾 Storage: free `{human_bytes(st_free)}` / total `{human_bytes(st_total)}`")
+    if ping is None:
+        lines.append("📡 Ping: `failed`")
+    else:
+        lines.append(f"📡 Ping (TCP): `{ping:.1f} ms`")
+
+    # Download
+    if dl.get("ok"):
+        lines.append(f"⬇️ Download: `{dl['mbps']:.2f} Mbps`  ({human_bytes(dl['bytes'])} in {dl['sec']:.2f}s)")
+    else:
+        lines.append(f"⬇️ Download: `failed` ({dl.get('err','unknown')})")
+
+    # Upload
+    if ul.get("ok"):
+        lines.append(f"⬆️ Upload: `{ul['mbps']:.2f} Mbps`  ({human_bytes(ul['bytes'])} in {ul['sec']:.2f}s)")
+    else:
+        lines.append(f"⬆️ Upload: `failed` ({ul.get('err','unknown')})")
+
+    # Worker2 KV RTT
+    if "w2_kv_rtt_ms" in r:
+        if r.get("w2_kv_rtt_ms") is None:
+            lines.append(f"🗄️ Worker2 KV RTT: `failed` ({r.get('w2_err','')})")
+        else:
+            lines.append(f"🗄️ Worker2 KV RTT: `{r['w2_kv_rtt_ms']:.1f} ms` (ok={r.get('w2_ok')})")
+
+    return "\n".join(lines)