From 5e83d66fd2be6b047e9ddf90b5f2528e68e0adfb Mon Sep 17 00:00:00 2001
From: Nytelife26 <xtylerjrx@gmail.com>
Date: Fri, 28 Jun 2024 00:04:53 +0100
Subject: [PATCH] refactor: reimplement megabit output

---
 cfspeedtest/__main__.py   |  4 ++--
 cfspeedtest/cloudflare.py | 42 +++++++++++++++++++++++++++------------
 2 files changed, 31 insertions(+), 15 deletions(-)

diff --git a/cfspeedtest/__main__.py b/cfspeedtest/__main__.py
index 34e6fdc..8985c13 100644
--- a/cfspeedtest/__main__.py
+++ b/cfspeedtest/__main__.py
@@ -4,7 +4,7 @@
 import sys
 from argparse import ArgumentParser
 
-from cfspeedtest.cloudflare import CloudflareSpeedtest
+from cfspeedtest.cloudflare import CloudflareSpeedtest, SuiteResults
 from cfspeedtest.logger import log, set_verbosity, setup_log
 from cfspeedtest.version import __version__
 
@@ -36,7 +36,7 @@ def cfspeedtest() -> None:
         log.debug("Python %s", sys.version)
         sys.exit(0)
 
-    results = CloudflareSpeedtest().run_all()
+    results = CloudflareSpeedtest(SuiteResults(megabits=not args.bps)).run_all()
 
     if args.json:
         setup_log()
diff --git a/cfspeedtest/cloudflare.py b/cfspeedtest/cloudflare.py
index 5b59760..56e9d82 100644
--- a/cfspeedtest/cloudflare.py
+++ b/cfspeedtest/cloudflare.py
@@ -137,6 +137,7 @@ def __init__(self, *, megabits: bool = False):
         super().__init__()
         self.setdefault("tests", {})
         self._megabits = megabits
+        log.info("megabits: %s", self._megabits)
 
     @property
     def meta(self) -> TestMetadata:
@@ -146,30 +147,44 @@ def meta(self) -> TestMetadata:
     def meta(self, value: TestMetadata) -> None:
         self["meta"] = value
         for meta_field, meta_value in value._asdict().items():
-            log.info("%s: %s", meta_field, meta_value)
+            log.info("meta.%s: %s", meta_field, meta_value)
 
     @property
     def tests(self) -> dict[str, TestResult]:
         return self["tests"]
 
-    def add_test(self, label: str, result: TestResult):
+    def add_test(self, label: str, result: TestResult) -> None:
+        if (
+            self._megabits
+            and result.value is not None
+            and label not in {"latency", "jitter"}
+        ):
+            result = TestResult(bits_to_megabits(result.value), result.time)
         self.tests[label] = result
+        log.info("tests.%s: %s", label, result.value)
+
+    def add_90th(self, test_type: TestType, result: TestResult) -> None:
+        if self._megabits and result.value is not None:
+            result = TestResult(bits_to_megabits(result.value), result.time)
+        label = f"90th_percentile_{test_type.name.lower()}"
+        self[label] = result
         log.info("%s: %s", label, result.value)
 
     @property
-    def percentile_90th_down_bps(self) -> TestResult:
-        return self["90th_percentile_down_bps"]
+    def percentile_90th_down(self) -> TestResult:
+        return self["90th_percentile_down"]
 
     @property
-    def percentile_90th_up_bps(self) -> TestResult:
-        return self["90th_percentile_up_bps"]
+    def percentile_90th_up(self) -> TestResult:
+        return self["90th_percentile_up"]
 
     def to_full_dict(self) -> dict:
         return {
+            "megabits": self._megabits,
             "meta": self.meta._asdict(),
             "tests": {k: v._asdict() for k, v in self.tests.items()},
-            "90th_percentile_down_bps": self.percentile_90th_down_bps,
-            "90th_percentile_up_bps": self.percentile_90th_up_bps,
+            "90th_percentile_down": self.percentile_90th_down,
+            "90th_percentile_up": self.percentile_90th_up,
         }
 
 
@@ -231,7 +246,8 @@ def run_test(self, test: TestSpec) -> TestTimers:
             )
             coll.full.append(time.time() - start)
             coll.server.append(
-                float(r.headers["Server-Timing"].split("=")[1].split(",")[0]) / 1e3
+                float(r.headers["Server-Timing"].split("=")[1].split(",")[0])
+                / 1e3
             )
             coll.request.append(
                 r.elapsed.seconds + r.elapsed.microseconds / 1e6
@@ -254,7 +270,7 @@ def run_test_speed(self, test: TestSpec) -> list[int]:
         """Run a test specification and collect speed results."""
         speeds = self.run_test(test).to_speeds(test)
         self.results.add_test(
-            f"{test.name}_{test.type.name.lower()}_bps",
+            f"{test.name}_{test.type.name.lower()}",
             TestResult(int(statistics.mean(speeds))),
         )
         return speeds
@@ -263,17 +279,17 @@ def run_all(self) -> SuiteResults:
         """Run the full test suite."""
         self.results.meta = self.metadata()
 
-        data = {"down": [], "up": []}
+        data = {TestType.Down: [], TestType.Up: []}
         for test in self.tests:
             if test.name == "latency":
                 self.run_test_latency(test)
                 continue
-            data[test.type.name.lower()].extend(self.run_test_speed(test))
+            data[test.type].extend(self.run_test_speed(test))
 
         for k, v in data.items():
             result = None
             if len(v) > 0:
                 result = int(_calculate_percentile(v, 0.9))
-            self.results[f"90th_percentile_{k}_bps"] = TestResult(result)
+            self.results.add_90th(k, TestResult(result))
 
         return self.results