Update bot/core/speedtest.py

bot/core/speedtest.py

@@ -1,170 +1,82 @@
 # PATH: bot/core/speedtest.py
-import asyncio
+import os
 import time
 import shutil
-from typing import Dict, Any, Optional
-
+import secrets
 import httpx
 
-from bot.core.uptime import uptime_seconds
-from bot.core.progress import human_bytes, human_eta
+# Defaults (can be overridden via env)
+# Cloudflare speed test endpoints (no "HF" mention anywhere)
+DL_URL = os.environ.get("SPEEDTEST_DL_URL", "https://speed.cloudflare.com/__down")
+UP_URL = os.environ.get("SPEEDTEST_UP_URL", "https://speed.cloudflare.com/__up")
+PING_URL = os.environ.get("SPEEDTEST_PING_URL", "https://www.google.com/generate_204")
+
+DL_BYTES_DEFAULT = int(float(os.environ.get("SPEEDTEST_DL_BYTES", "8000000")))  # ~8MB
+UP_BYTES_DEFAULT = int(float(os.environ.get("SPEEDTEST_UP_BYTES", "3000000")))  # ~3MB
+TIMEOUT_SEC = float(os.environ.get("SPEEDTEST_TIMEOUT_SEC", "30"))
 
-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 bytes_to_mb(x: float) -> float:
+    # "real mb/s" => treat 1 MB = 1024*1024 bytes
+    return float(x) / (1024.0 * 1024.0)
 
-def _bps_to_mbps(bps: float) -> float:
-    return (bps * 8.0) / 1_000_000.0
+def bytes_per_sec_to_mb_s(bps: float) -> float:
+    return bytes_to_mb(bps)
+
+def disk_total_free(path: str = "/") -> dict:
+    du = shutil.disk_usage(path)
+    return {"total": du.total, "free": du.free, "used": du.used}
 
-async def tcp_ping_ms(host: str = "1.1.1.1", port: int = 443, timeout: float = 3.0) -> Optional[float]:
+async def ping_ms() -> float | None:
     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
+        async with httpx.AsyncClient(timeout=TIMEOUT_SEC, follow_redirects=True) as c:
+            r = await c.get(PING_URL, headers={"cache-control": "no-cache"})
+            # Even if status not 204, RTT still valid for "ping-like"
+            _ = r.status_code
         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)
+async def net_download_test(bytes_to_get: int | None = None) -> dict:
+    n = int(bytes_to_get or DL_BYTES_DEFAULT)
     t0 = time.perf_counter()
-    total = 0
+    got = 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),
-    }
+    try:
+        async with httpx.AsyncClient(timeout=TIMEOUT_SEC, follow_redirects=True) as c:
+            # Cloudflare accepts ?bytes=...
+            url = f"{DL_URL}?bytes={n}"
+            async with c.stream("GET", url, headers={"cache-control": "no-cache"}) as r:
+                r.raise_for_status()
+                async for chunk in r.aiter_bytes():
+                    got += len(chunk)
+                    if got >= n:
+                        break
+    except Exception:
+        # return partial if any
+        pass
 
-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
+    sec = max(0.001, time.perf_counter() - t0)
+    bps = float(got) / sec
+    return {"bytes": got, "seconds": sec, "bps": bps}
 
-    async def gen():
-        nonlocal remaining
-        while remaining > 0:
-            take = chunk if remaining >= len(chunk) else b"\x00" * remaining
-            remaining -= len(take)
-            yield take
+async def net_upload_test(bytes_to_send: int | None = None) -> dict:
+    n = int(bytes_to_send or UP_BYTES_DEFAULT)
+    payload = secrets.token_bytes(n)
 
     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
+    sent = 0
     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')})")
+        async with httpx.AsyncClient(timeout=TIMEOUT_SEC, follow_redirects=True) as c:
+            # Cloudflare accepts ?bytes=... (and body upload)
+            url = f"{UP_URL}?bytes={n}"
+            r = await c.post(url, content=payload, headers={"content-type": "application/octet-stream"})
+            # If server returns error, still count attempt
+            _ = r.status_code
+        sent = n
+    except Exception:
+        sent = 0
 
-    return "\n".join(lines)
+    sec = max(0.001, time.perf_counter() - t0)
+    bps = float(sent) / sec
+    return {"bytes": sent, "seconds": sec, "bps": bps}