Improve speedtest output, bump version

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "TrafficToll"
-version = "1.3.1"
+version = "1.4.0"
 description = "NetLimiter-like bandwidth limiting and QoS for Linux"
 authors = ["Chris Braun <[email protected]>"]
 

traffictoll/__main__.py

@@ -3,7 +3,7 @@
 from loguru import logger
 
 from .cli import get_argument_parser, main as cli_main
-from .exceptions import ConfigError, DependencyError
+from .exceptions import ConfigError, MissingDependencyError
 
 
 def main() -> None:
@@ -19,7 +19,7 @@ def main() -> None:
         logger.info("Aborted")
     except ConfigError as error:
         logger.error("Invalid configuration: {}", error)
-    except DependencyError as error:
+    except MissingDependencyError as error:
         logger.error("Missing dependency: {}", error)
     except Exception:
         logger.exception("Unexpected error occurred:")

traffictoll/cli.py

@@ -8,6 +8,7 @@
 from loguru import logger
 from ruamel.yaml import YAML
 
+from .exceptions import DependencyOutputError, MissingDependencyError
 from .net import ProcessFilterPredicate, filter_net_connections
 from .speedtest import test_speed
 from .tc import (
@@ -82,12 +83,21 @@ def main(arguments: argparse.Namespace) -> None:
     config_global_upload_rate = config.get("upload")
     if arguments.speed_test:
         logger.info("Running speed test...")
-        results = test_speed()
-        if results:
+
+        try:
+            result = test_speed()
+        except MissingDependencyError as error:
+            logger.error("Missing dependency: {}", error)
+            result = None
+        except DependencyOutputError as error:
+            logger.error("Dependency output error: {}", error)
+            result = None
+
+        if result:
             logger.info(
-                "Determined download speed: {}bps, upload speed: {}bps", *results
+                "Determined download speed: {}bps, upload speed: {}bps", *result
             )
-            config_global_download_rate, config_global_upload_rate = results
+            config_global_download_rate, config_global_upload_rate = result
         else:
             logger.error(
                 "Failed to automatically determine download and upload speed, falling "

traffictoll/exceptions.py

@@ -1,6 +1,14 @@
-class DependencyError(Exception):
+class TrafficTollException(Exception):
     pass
 
 
-class ConfigError(Exception):
+class MissingDependencyError(TrafficTollException):
+    pass
+
+
+class DependencyOutputError(TrafficTollException):
+    pass
+
+
+class ConfigError(TrafficTollException):
     pass

traffictoll/speedtest.py

@@ -1,49 +1,74 @@
+import collections
+import enum
 import json
 import subprocess
-from typing import Tuple, Optional
+from typing import Optional
 
-from .exceptions import DependencyError
+from .exceptions import DependencyOutputError
 from .utils import run
 
+SPEEDTEST_VERSION_COMMAND = "speedtest --version"
+OOKLA_SPEEDTEST_COMMAND = "speedtest --format=json"
+SIVEL_SPEEDTEST_COMMAND = "speedtest --json"
+
+SpeedTestResult = collections.namedtuple("SpeedTest", ["download_rate", "upload_rate"])
+
+
+class _SpeedTestProvider(enum.Enum):
+    Ookla = enum.auto()
+    Sivel = enum.auto()
+    Unknown = enum.auto()
+
 
 # https://www.speedtest.net/apps/cli
-def _ookla_speedtest_cli() -> Optional[Tuple[int, int]]:
+def _ookla_speedtest_cli() -> SpeedTestResult:
     process = run(
-        "speedtest --format=json", stdout=subprocess.PIPE, universal_newlines=True,
+        OOKLA_SPEEDTEST_COMMAND, stdout=subprocess.PIPE, universal_newlines=True,
     )
 
     try:
         result = json.loads(process.stdout)
-        return result["download"]["bandwidth"], result["upload"]["bandwidth"]
+        return SpeedTestResult(
+            result["download"]["bandwidth"], result["upload"]["bandwidth"]
+        )
     except (json.JSONDecodeError, KeyError):
-        return None
+        raise DependencyOutputError(
+            f"Command: {OOKLA_SPEEDTEST_COMMAND!r} returned unrecognized output: "
+            f"{process.stdout!r}"
+        )
 
 
 # https://github.com/sivel/speedtest-cli
-def _sivel_speedtest_cli() -> Optional[Tuple[int, int]]:
-    process = run("speedtest --json", stdout=subprocess.PIPE, universal_newlines=True)
+def _sivel_speedtest_cli() -> SpeedTestResult:
+    process = run(
+        SIVEL_SPEEDTEST_COMMAND, stdout=subprocess.PIPE, universal_newlines=True
+    )
 
     try:
         result = json.loads(process.stdout)
-        return round(result["download"]), round(result["upload"])
+        return SpeedTestResult(round(result["download"]), round(result["upload"]))
     except (json.JSONDecodeError, KeyError):
-        pass
+        raise DependencyOutputError(
+            f"Command: {SIVEL_SPEEDTEST_COMMAND!r} returned unrecognized output: "
+            f"{process.stdout!r}"
+        )
 
 
-def test_speed() -> Optional[Tuple[int, int]]:
-    try:
-        process = run(
-            "speedtest --version", stdout=subprocess.PIPE, universal_newlines=True
-        )
-    except DependencyError:
-        return
+def _get_speedtest_provider() -> _SpeedTestProvider:
+    process = run(
+        SPEEDTEST_VERSION_COMMAND, stdout=subprocess.PIPE, universal_newlines=True
+    )
+    if process.stdout.startswith("Speedtest by Ookla"):
+        return _SpeedTestProvider.Ookla
+    elif process.stdout.startswith("speedtest-cli"):
+        return _SpeedTestProvider.Sivel
+    return _SpeedTestProvider.Unknown
+
 
-    lines = process.stdout.splitlines()
-    if not lines:
-        return
+def test_speed() -> Optional[SpeedTestResult]:
+    speedtest_version = _get_speedtest_provider()
 
-    first_line = process.stdout.splitlines()[0]
-    if first_line.startswith("Speedtest by Ookla"):
+    if speedtest_version is _SpeedTestProvider.Ookla:
         return _ookla_speedtest_cli()
-    elif first_line.startswith("speedtest-cli"):
+    elif speedtest_version is _SpeedTestProvider.Sivel:
         return _sivel_speedtest_cli()

traffictoll/utils.py

@@ -6,7 +6,7 @@
 
 from loguru import logger
 
-from .exceptions import DependencyError
+from .exceptions import MissingDependencyError
 
 
 # Cache the full executable path just in case subprocess.Popen doesn't
@@ -19,7 +19,7 @@ def run(command: str, **popen_kwargs) -> subprocess.CompletedProcess:
     executable, *arguments = shlex.split(command)
     path = _which(executable)
     if not path:
-        raise DependencyError(f"Executable for command: {command!r} not found")
+        raise MissingDependencyError(f"Executable for command: {command!r} not found")
 
     logger.debug(command)
     return subprocess.run([path] + arguments, **popen_kwargs)