Test sockio and make its docs better

sbws/commands/client.py

@@ -17,6 +17,7 @@ from argparse import ArgumentDefaultsHelpFormatter
 from multiprocessing.dummy import Pool
 from threading import Event
 from threading import RLock
+import traceback
 import socks
 import socket
 import time
@@ -100,7 +101,7 @@ def measure_rtt_to_server(sock, conf):
     rtts = []
     for _ in range(0, conf.getint('client', 'num_rtts')):
         start_time = time.time()
-        if not tell_server_amount(sock, MIN_REQ_BYTES):
+        if not tell_server_amount(sock, -1):
             log.info('Unable to ping server on', sock.fileno())
             return
         try:
@@ -287,6 +288,7 @@ def result_putter_error(target):
     measurement -- and return that function so it can be used by someone else
     '''
     def closure(err):
+        log.warn(traceback.format_exc())
         log.warn('Unhandled exception caught while measuring {}: {} {}'.format(
             target.nickname, type(err), err))
     return closure
@@ -316,7 +318,7 @@ def test_speedtest(args, conf):
             callback_err = result_putter_error(target)
             async_result = pool.apply_async(
                 measure_relay, [args, conf, helpers, cb, rl, target], {},
-                callback, callback_err)
+                callback)
             pending_results.append(async_result)
             while len(pending_results) >= max_pending_results:
                 time.sleep(5)

sbws/util/sockio.py

@@ -2,15 +2,28 @@ import socket
 
 
 def read_line(s, max_len=None, log_fn=print):
-    ''' read until b'\n' is seen on the socket <s>. Return everything up until
-    the newline as a str. If nothing can be read, return None. Note how that is
-    different than if a newline is the first character; in that case, an empty
-    str is returned.
+    '''
+    Read from the blocking socket **s** until nothing can be read anymore or
+    until a b'\n' is seen.
 
     If max_len is specified, then that's the maximum number of characters that
     will be returned, even if a newline is not read yet.
+
+    This function can only handle characters that can be represented as a
+    single byte in utf8.
+
+    :param socket.socket s: Blocking socket to read from
+    :param int max_len: Maximum number of bytes to read, not including a
+        tailing newline
+    :param func log_fn: Function with a signature similar to ``print`` to call
+        if an error occurs
+    :raises UnicodeDecodeError: If any byte is not a valid utf8 byte
+    :returns: Everything read up until a newline and with a maximum length of
+        **max_len**. If nothing could be read, returns None. If a newline is
+        the first character, returns an empty string.
     '''
-    assert max_len is None or max_len > 0
+    assert isinstance(s, socket.socket)
+    assert max_len is None or (isinstance(max_len, int) and max_len > 0)
     chars = None
     while True:
         try:
@@ -24,7 +37,10 @@ def read_line(s, max_len=None, log_fn=print):
             chars = ''
         if c == b'\n':
             break
-        chars += c.decode('utf-8')
+        try:
+            chars += c.decode('utf-8')
+        except UnicodeDecodeError as e:
+            raise e
         if max_len is not None and len(chars) >= max_len:
             return chars[0:max_len]
     return chars

tests/util/test_sockio.py

+from sbws.util.sockio import read_line
+import socket
+
+
+class MockReadingSocket(socket.socket):
+    def __init__(self, data):
+        assert isinstance(data, bytes)
+        self._data = data
+        self._next = 0
+
+    def recv(self, amount):
+        assert amount == 1, 'read_line should only ever request one byte'
+        start = self._next
+        end = start + amount
+        ret = self._data[start:end]
+        self._next += len(ret)
+        return ret
+
+
+class MockReadingSocketWithTimeout(socket.socket):
+    def __init__(self):
+        pass
+
+    def recv(self, amount):
+        raise socket.timeout()
+
+
+def test_sockio_simple():
+    in_str = b'1234\n'
+    expected = '1234'
+    sock = MockReadingSocket(in_str)
+    out = read_line(sock)
+    assert out == expected
+
+
+def test_sockio_empty():
+    in_str = b'\n'
+    expected = ''
+    sock = MockReadingSocket(in_str)
+    out = read_line(sock)
+    assert out == expected
+
+
+def test_sockio_null():
+    in_str = b''
+    expected = None
+    sock = MockReadingSocket(in_str)
+    out = read_line(sock)
+    assert out == expected
+
+
+def test_sockio_too_much():
+    in_str = b'12345678\n'
+    expected = '1234'
+    sock = MockReadingSocket(in_str)
+    out = read_line(sock, max_len=len(expected))
+    assert out == expected
+
+
+def test_sockio_timeout():
+    expected = None
+    sock = MockReadingSocketWithTimeout()
+    try:
+        out = read_line(sock)
+    except socket.timeout:
+        assert None, 'Should not have let the timeout bubble up'
+    assert out == expected
+
+
+def test_sockio_missing_newline():
+    in_str = b'1234'
+    expected = '1234'
+    sock = MockReadingSocket(in_str)
+    out = read_line(sock)
+    assert out == expected
+
+
+def test_sockio_bad_max_len():
+    in_str = b'1234'
+    sock = MockReadingSocket(in_str)
+    try:
+        read_line(sock, max_len=1.2)
+    except AssertionError:
+        pass
+    else:
+        assert None, 'Should have failed'
+
+
+def test_sockio_non_ascii():
+    in_str = b'asdf\x80asdf'
+    sock = MockReadingSocket(in_str)
+    try:
+        read_line(sock)
+    except UnicodeDecodeError:
+        pass
+    else:
+        assert None, 'Should have failed'