Update bot/core/speedtest.py

bot/core/speedtest.py

@@ -1,82 +1,82 @@
 # PATH: bot/core/speedtest.py
+import asyncio
 import os
 import time
-import shutil
-import secrets
+from typing import Optional
+
 import httpx
 
-# 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")
+CF_DOWN = "https://speed.cloudflare.com/__down"
+CF_UP = "https://speed.cloudflare.com/__up"
+
 
-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"))
+def bytes_to_mb(n: float) -> float:
+    return float(n) / (1024.0 * 1024.0)
 
-def bytes_to_mb(x: float) -> float:
-    # "real mb/s" => treat 1 MB = 1024*1024 bytes
-    return float(x) / (1024.0 * 1024.0)
 
 def bytes_per_sec_to_mb_s(bps: float) -> float:
-    return bytes_to_mb(bps)
+    return float(bps) / (1024.0 * 1024.0)
+
 
 def disk_total_free(path: str = "/") -> dict:
-    du = shutil.disk_usage(path)
-    return {"total": du.total, "free": du.free, "used": du.used}
+    st = os.statvfs(path)
+    total = st.f_frsize * st.f_blocks
+    free = st.f_frsize * st.f_bavail
+    return {"total": int(total), "free": int(free)}
 
-async def ping_ms() -> float | None:
+
+async def ping_ms(host: str = "1.1.1.1", port: int = 443, timeout: float = 3.0) -> Optional[float]:
+    """
+    "Ping" approximation via TCP connect time (no ICMP).
+    Returns ms or None.
+    """
     t0 = time.perf_counter()
     try:
-        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
+        fut = asyncio.open_connection(host, port)
+        reader, writer = await asyncio.wait_for(fut, timeout=timeout)
+        try:
+            writer.close()
+            if hasattr(writer, "wait_closed"):
+                await writer.wait_closed()
+        except Exception:
+            pass
         return (time.perf_counter() - t0) * 1000.0
     except Exception:
         return None
 
-async def net_download_test(bytes_to_get: int | None = None) -> dict:
-    n = int(bytes_to_get or DL_BYTES_DEFAULT)
+
+async def net_download_test(bytes_target: int = 8 * 1024 * 1024, timeout: float = 15.0) -> dict:
+    """
+    Downloads bytes_target from Cloudflare endpoint and measures throughput.
+    """
+    url = f"{CF_DOWN}?bytes={int(bytes_target)}"
     t0 = time.perf_counter()
     got = 0
 
-    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
+    limits = httpx.Limits(max_connections=10, max_keepalive_connections=5)
+    async with httpx.AsyncClient(timeout=timeout, limits=limits, follow_redirects=True) as client:
+        async with client.stream("GET", url) as r:
+            r.raise_for_status()
+            async for chunk in r.aiter_bytes():
+                got += len(chunk)
 
-    sec = max(0.001, time.perf_counter() - t0)
-    bps = float(got) / sec
-    return {"bytes": got, "seconds": sec, "bps": bps}
+    sec = max(1e-6, time.perf_counter() - t0)
+    bps = got / sec
+    return {"bytes": int(got), "seconds": float(sec), "bps": float(bps)}
 
-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)
 
+async def net_upload_test(bytes_target: int = 3 * 1024 * 1024, timeout: float = 20.0) -> dict:
+    """
+    Uploads bytes_target to Cloudflare endpoint and measures throughput.
+    """
+    payload = os.urandom(int(bytes_target))
     t0 = time.perf_counter()
-    sent = 0
-    try:
-        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
 
-    sec = max(0.001, time.perf_counter() - t0)
-    bps = float(sent) / sec
-    return {"bytes": sent, "seconds": sec, "bps": bps}
+    limits = httpx.Limits(max_connections=10, max_keepalive_connections=5)
+    async with httpx.AsyncClient(timeout=timeout, limits=limits, follow_redirects=True) as client:
+        r = await client.post(CF_UP, content=payload, headers={"content-type": "application/octet-stream"})
+        r.raise_for_status()
+
+    sec = max(1e-6, time.perf_counter() - t0)
+    bps = bytes_target / sec
+    return {"bytes": int(bytes_target), "seconds": float(sec), "bps": float(bps)}