First pass at adding ability to test using newer socket based connection

1 files changed, 197 insertions, 55 deletions

diff --git a/speedtest.py b/speedtest.py
index 0fe067b..5e6a6f5 100755
--- a/speedtest.py
+++ b/speedtest.py
@@ -713,7 +713,7 @@ class HTTPDownloader(threading.Thread):
                  shutdown_event=None):
         threading.Thread.__init__(self)
         self.request = request
-        self.result = [0]
+        self.result = 0
         self.starttime = start
         self.timeout = timeout
         self.i = i
@@ -734,14 +734,62 @@ class HTTPDownloader(threading.Thread):
                 while (not self._shutdown_event.isSet() and
                         (timeit.default_timer() - self.starttime) <=
                         self.timeout):
-                    self.result.append(len(f.read(10240)))
-                    if self.result[-1] == 0:
+                    data = len(f.read(10240))
+                    if data == 0:
                         break
+                    self.result += data
                 f.close()
         except IOError:
             pass
 
 
+class SocketDownloader(threading.Thread):
+    def __init__(self, i, address, size, start, timeout, shutdown_event=None,
+                 source_address=None):
+        threading.Thread.__init__(self)
+        self.result = 0
+        self.starttime = start
+        self.timeout = timeout
+        self.i = i
+        self.size = size
+        self.remaining = self.size
+
+        if shutdown_event:
+            self._shutdown_event = shutdown_event
+        else:
+            self._shutdown_event = FakeShutdownEvent()
+
+        self.sock = create_connection(address, timeout=timeout,
+                                      source_address=source_address)
+
+    def run(self):
+        try:
+            if (timeit.default_timer() - self.starttime) <= self.timeout:
+                self.sock.sendall('HI\n'.encode())
+                self.sock.recv(1024)
+
+                while (self.remaining and not self._shutdown_event.isSet() and
+                        (timeit.default_timer() - self.starttime) <=
+                        self.timeout):
+
+                    if self.remaining > 1000000:
+                        ask = 1000000
+                    else:
+                        ask = self.remaining
+
+                    down = 0
+                    self.sock.sendall(('DOWNLOAD %d\n' % ask).encode())
+                    while down < ask:
+                        down += len(self.sock.recv(10240))
+
+                    self.result += down
+                    self.remaining -= down
+
+                self.sock.close()
+        except IOError:
+            pass
+
+
 class HTTPUploaderData(object):
     """File like object to improve cutting off the upload once the timeout
     has been reached
@@ -759,7 +807,7 @@ class HTTPUploaderData(object):
 
         self._data = None
 
-        self.total = [0]
+        self.total = 0
 
     def pre_allocate(self):
         chars = '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ'
@@ -787,7 +835,7 @@ class HTTPUploaderData(object):
         if ((timeit.default_timer() - self.start) <= self.timeout and
                 not self._shutdown_event.isSet()):
             chunk = self.data.read(n)
-            self.total.append(len(chunk))
+            self.total += len(chunk)
             return chunk
         else:
             raise SpeedtestUploadTimeout()
@@ -835,11 +883,59 @@ class HTTPUploader(threading.Thread):
                     f = self._opener(request)
                 f.read(11)
                 f.close()
-                self.result = sum(self.request.data.total)
+                self.result = self.request.data.total
             else:
                 self.result = 0
         except (IOError, SpeedtestUploadTimeout):
-            self.result = sum(self.request.data.total)
+            self.result = self.request.data.total
+
+
+class SocketUploader(threading.Thread):
+    def __init__(self, i, address, size, start, timeout, shutdown_event=None,
+                 source_address=None):
+        threading.Thread.__init__(self)
+        self.result = 0
+        self.starttime = start
+        self.timeout = timeout
+        self.i = i
+        self.size = size
+        self.remaining = self.size
+
+        if shutdown_event:
+            self._shutdown_event = shutdown_event
+        else:
+            self._shutdown_event = FakeShutdownEvent()
+
+        self.sock = create_connection(address, timeout=timeout,
+                                      source_address=source_address)
+
+    def run(self):
+        try:
+            if (timeit.default_timer() - self.starttime) <= self.timeout:
+                self.sock.sendall('HI\n'.encode())
+                self.sock.recv(1024)
+
+                while (self.remaining and not self._shutdown_event.isSet() and
+                        (timeit.default_timer() - self.starttime) <=
+                        self.timeout):
+
+                    if self.remaining > 100000:
+                        give = 100000
+                    else:
+                        give = self.remaining
+
+                    header = ('UPLOAD %d 0\n' % give).encode()
+                    data = '0'.encode() * (give - len(header))
+
+                    self.sock.sendall(header)
+                    self.sock.sendall(data)
+                    self.sock.recv(24)
+                    self.result += give
+                    self.remaining -= give
+
+                self.sock.close()
+        except IOError:
+            pass
 
 
 class SpeedtestResults(object):
@@ -997,7 +1093,7 @@ class Speedtest(object):
     """Class for performing standard speedtest.net testing operations"""
 
     def __init__(self, config=None, source_address=None, timeout=10,
-                 secure=False, shutdown_event=None):
+                 secure=False, shutdown_event=None, use_socket=False):
         self.config = {}
 
         self._source_address = source_address
@@ -1011,6 +1107,8 @@ class Speedtest(object):
         else:
             self._shutdown_event = FakeShutdownEvent()
 
+        self._use_socket = use_socket
+
         self.get_config()
         if config is not None:
             self.config.update(config)
@@ -1105,9 +1203,14 @@ class Speedtest(object):
         up_sizes = [32768, 65536, 131072, 262144, 524288, 1048576, 7340032]
         sizes = {
             'upload': up_sizes[ratio - 1:],
-            'download': [350, 500, 750, 1000, 1500, 2000, 2500,
-                         3000, 3500, 4000]
         }
+        if self._use_socket:
+            sizes['download'] = [245388, 505544, 1118012, 1986284, 4468241,
+                                 7907740, 12407926, 17816816, 24262167,
+                                 31625365]
+        else:
+            sizes['download'] = [350, 500, 750, 1000, 1500, 2000, 2500,
+                                 3000, 3500, 4000]
 
         size_count = len(sizes['upload'])
 
@@ -1252,6 +1355,9 @@ class Speedtest(object):
                             or int(attrib.get('id')) in exclude):
                         continue
 
+                    host, port = attrib['host'].split(':')
+                    attrib['host'] = (host, int(port))
+
                     try:
                         d = distance(self.lat_lon,
                                      (float(attrib.get('lat')),
@@ -1429,29 +1535,48 @@ class Speedtest(object):
     def download(self, callback=do_nothing):
         """Test download speed against speedtest.net"""
 
-        urls = []
-        for size in self.config['sizes']['download']:
-            for _ in range(0, self.config['counts']['download']):
-                urls.append('%s/random%sx%s.jpg' %
-                            (os.path.dirname(self.best['url']), size, size))
+        if self._use_socket:
+            requests = []
+            for size in self.config['sizes']['download']:
+                for _ in range(0, self.config['counts']['download']):
+                    requests.append(size)
 
-        request_count = len(urls)
-        requests = []
-        for i, url in enumerate(urls):
-            requests.append(
-                build_request(url, bump=i, secure=self._secure)
-            )
+            request_count = len(requests)
+        else:
+            urls = []
+            for size in self.config['sizes']['download']:
+                for _ in range(0, self.config['counts']['download']):
+                    urls.append('%s/random%sx%s.jpg' %
+                                (os.path.dirname(self.best['url']), size, size))
+
+            request_count = len(urls)
+            requests = []
+            for i, url in enumerate(urls):
+                requests.append(
+                    build_request(url, bump=i, secure=self._secure)
+                )
 
         def producer(q, requests, request_count):
             for i, request in enumerate(requests):
-                thread = HTTPDownloader(
-                    i,
-                    request,
-                    start,
-                    self.config['length']['download'],
-                    opener=self._opener,
-                    shutdown_event=self._shutdown_event
-                )
+                if self._use_socket:
+                    thread = SocketDownloader(
+                        i,
+                        self.best['host'],
+                        request,
+                        start,
+                        self.config['length']['download'],
+                        shutdown_event=self._shutdown_event,
+                        source_address=self._source_address
+                    )
+                else:
+                    thread = HTTPDownloader(
+                        i,
+                        request,
+                        start,
+                        self.config['length']['download'],
+                        opener=self._opener,
+                        shutdown_event=self._shutdown_event
+                    )
                 thread.start()
                 q.put(thread, True)
                 callback(i, request_count, start=True)
@@ -1463,7 +1588,7 @@ class Speedtest(object):
                 thread = q.get(True)
                 while thread.isAlive():
                     thread.join(timeout=0.1)
-                finished.append(sum(thread.result))
+                finished.append(thread.result)
                 callback(thread.i, request_count, end=True)
 
         q = Queue(self.config['threads']['download'])
@@ -1502,34 +1627,48 @@ class Speedtest(object):
 
         requests = []
         for i, size in enumerate(sizes):
-            # We set ``0`` for ``start`` and handle setting the actual
-            # ``start`` in ``HTTPUploader`` to get better measurements
-            data = HTTPUploaderData(
-                size,
-                0,
-                self.config['length']['upload'],
-                shutdown_event=self._shutdown_event
-            )
-            if pre_allocate:
-                data.pre_allocate()
-            requests.append(
-                (
-                    build_request(self.best['url'], data, secure=self._secure),
-                    size
+            if self._use_socket:
+                requests.append(size)
+            else:
+                # We set ``0`` for ``start`` and handle setting the actual
+                # ``start`` in ``HTTPUploader`` to get better measurements
+                data = HTTPUploaderData(
+                    size,
+                    0,
+                    self.config['length']['upload'],
+                    shutdown_event=self._shutdown_event
+                )
+                if pre_allocate:
+                    data.pre_allocate()
+                requests.append(
+                    (
+                        build_request(self.best['url'], data, secure=self._secure),
+                        size
+                    )
                 )
-            )
 
         def producer(q, requests, request_count):
             for i, request in enumerate(requests[:request_count]):
-                thread = HTTPUploader(
-                    i,
-                    request[0],
-                    start,
-                    request[1],
-                    self.config['length']['upload'],
-                    opener=self._opener,
-                    shutdown_event=self._shutdown_event
-                )
+                if self._use_socket:
+                    thread = SocketUploader(
+                        i,
+                        self.best['host'],
+                        request,
+                        start,
+                        self.config['length']['upload'],
+                        shutdown_event=self._shutdown_event,
+                        source_address=self._source_address
+                    )
+                else:
+                    thread = HTTPUploader(
+                        i,
+                        request[0],
+                        start,
+                        request[1],
+                        self.config['length']['upload'],
+                        opener=self._opener,
+                        shutdown_event=self._shutdown_event
+                    )
                 thread.start()
                 q.put(thread, True)
                 callback(i, request_count, start=True)
@@ -1652,6 +1791,8 @@ def parse_args():
     parser.add_argument('--secure', action='store_true',
                         help='Use HTTPS instead of HTTP when communicating '
                              'with speedtest.net operated servers')
+    parser.add_argument('--socket', action='store_true',
+                        help='Use socket test instead of HTTP based tests')
     parser.add_argument('--no-pre-allocate', dest='pre_allocate',
                         action='store_const', default=True, const=False,
                         help='Do not pre allocate upload data. Pre allocation '
@@ -1764,7 +1905,8 @@ def shell():
         speedtest = Speedtest(
             source_address=args.source,
             timeout=args.timeout,
-            secure=args.secure
+            secure=args.secure,
+            use_socket=args.socket
         )
     except (ConfigRetrievalError,) + HTTP_ERRORS:
         printer('Cannot retrieve speedtest configuration', error=True)