From 506678a2135b69d4ab11203817a21584b6191b08 Mon Sep 17 00:00:00 2001
From: markafarrell <mark.andrew.farrell@gmail.com>
Date: Sun, 25 Oct 2020 21:09:39 +1100
Subject: [PATCH] add prometheus pushgateway option

---
 speedtest-cli.1 |  5 +++++
 speedtest.py    | 44 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+)

diff --git a/speedtest-cli.1 b/speedtest-cli.1
index 367c977..865efce 100644
--- a/speedtest-cli.1
+++ b/speedtest-cli.1
@@ -118,6 +118,11 @@ Do not pre allocate upload data. Pre allocation is enabled by default to improve
 Show the version number and exit
 .RE
 
+\fB\-\-prometheus-push-gateway\fR
+.RS
+URL of Prometheus PushGateway to push results to. e.g http://localhost:9021 or https://localhost:9021
+.RE
+
 .SH EXAMPLES
 
 \fBAutomatically find closest server and start testing\fR
diff --git a/speedtest.py b/speedtest.py
index 92a2be0..b9456e4 100755
--- a/speedtest.py
+++ b/speedtest.py
@@ -36,6 +36,12 @@ except ImportError:
     gzip = None
     GZIP_BASE = object
 
+try:
+    from prometheus_client import CollectorRegistry, Gauge, push_to_gateway
+    prometheus_client = True
+except ImportError:
+    prometheus_client = None
+
 __version__ = '2.1.2'
 
 
@@ -782,6 +788,15 @@ def print_dots(shutdown_event):
 def do_nothing(*args, **kwargs):
     pass
 
+class PrometheusPushGateway(object):
+    """ Class for managing prometheus push gateway configuration"""
+    def __init__(self, server):
+        self.server = server
+        self.registry = CollectorRegistry()
+        self.gauges = {}
+        self.gauges['download'] = Gauge('speedtest_dl_throughput_bps', 'Speedtest.Net Download Speed (bps)', registry=self.registry)
+        self.gauges['upload'] = Gauge('speedtest_ul_throughput_bps', 'Speedtest.Net Upload Speed (bps)', registry=self.registry)
+        self.gauges['latency'] = Gauge('speedtest_latency_ms', 'Speedtest.Net Latency (ms)', registry=self.registry)
 
 class HTTPDownloader(threading.Thread):
     """Thread class for retrieving a URL"""
@@ -1771,6 +1786,8 @@ def parse_args():
                         help='Show the version number and exit')
     parser.add_argument('--debug', action='store_true',
                         help=ARG_SUPPRESS, default=ARG_SUPPRESS)
+    parser.add_argument('--prometheus-push-gateway', dest='prometheus_push_gateway',
+                        help='URL of Prometheus PushGateway to push results to. e.g http://localhost:9091 or https://localhost:9091')
 
     options = parser.parse_args()
     if isinstance(options, tuple):
@@ -1790,6 +1807,7 @@ def validate_optional_args(args):
     optional_args = {
         'json': ('json/simplejson python module', json),
         'secure': ('SSL support', HTTPSConnection),
+        'prometheus_push_gateway': ('prometheus_client python module', prometheus_client)
     }
 
     for arg, info in optional_args.items():
@@ -1818,6 +1836,12 @@ def printer(string, quiet=False, debug=False, error=False, **kwargs):
     if not quiet:
         print_(out, **kwargs)
 
+def is_url(url):
+  try:
+    result = urlparse(url)
+    return all([result.scheme, result.netloc])
+  except ValueError:
+    return False
 
 def shell():
     """Run the full speedtest.net test"""
@@ -1837,6 +1861,10 @@ def shell():
         raise SpeedtestCLIError('Cannot supply both --no-download and '
                                 '--no-upload')
 
+    if args.prometheus_push_gateway and \
+            not is_url(args.prometheus_push_gateway):
+        raise SpeedtestCLIError('--prometheus-push-gateway must be a valid URL')
+
     if len(args.csv_delimiter) != 1:
         raise SpeedtestCLIError('--csv-delimiter must be a single character')
 
@@ -1980,6 +2008,22 @@ def shell():
     if args.share and not machine_format:
         printer('Share results: %s' % results.share())
 
+    if args.prometheus_push_gateway:
+        prometheus_push_gateway = PrometheusPushGateway(args.prometheus_push_gateway)
+
+        prometheus_push_gateway.gauges['upload'].set(results.upload)
+        prometheus_push_gateway.gauges['download'].set(results.download)
+        prometheus_push_gateway.gauges['latency'].set(results.ping)
+        try:
+            push_to_gateway(
+                prometheus_push_gateway.server,
+                job='speedtest',
+                registry=prometheus_push_gateway.registry,
+                timeout=None
+            )
+        except URLError:
+            # Couldn't talk to push gateway
+            printer('Could not contact Prometheus PushGateway', quiet)
 
 def main():
     try: