Better handling of retrieving config and server list

This commit is contained in:
Matt Martz 2017-01-09 19:48:35 -06:00
parent 411f1609e8
commit 823d7dc2b7
1 changed files with 15 additions and 9 deletions

View File

@ -297,7 +297,13 @@ class GzipDecodedResponse(GZIP_BASE):
raise SpeedtestHTTPError('HTTP response body is gzip encoded, ' raise SpeedtestHTTPError('HTTP response body is gzip encoded, '
'but gzip support is not available') 'but gzip support is not available')
IO = BytesIO or StringIO IO = BytesIO or StringIO
self.io = IO(response.read()) self.io = IO()
while 1:
chunk = response.read(1024)
if len(chunk) == 0:
break
self.io.write(chunk)
self.io.seek(0)
gzip.GzipFile.__init__(self, mode='rb', fileobj=self.io) gzip.GzipFile.__init__(self, mode='rb', fileobj=self.io)
def close(self): def close(self):
@ -721,7 +727,7 @@ class Speedtest(object):
stream = get_response_stream(uh) stream = get_response_stream(uh)
while 1: while 1:
configxml.append(stream.read(10240)) configxml.append(stream.read(1024))
if len(configxml[-1]) == 0: if len(configxml[-1]) == 0:
break break
stream.close() stream.close()
@ -833,7 +839,7 @@ class Speedtest(object):
serversxml = [] serversxml = []
while 1: while 1:
serversxml.append(stream.read(10240)) serversxml.append(stream.read(1024))
if len(serversxml[-1]) == 0: if len(serversxml[-1]) == 0:
break break
@ -1321,16 +1327,16 @@ def shell():
printer('Retrieving speedtest.net configuration...', quiet) printer('Retrieving speedtest.net configuration...', quiet)
try: try:
speedtest = Speedtest() speedtest = Speedtest()
except ConfigRetrievalError: except (ConfigRetrievalError, HTTP_ERRORS):
printer('Cannot retrieve speedtest configuration') printer('Cannot retrieve speedtest configuration')
raise raise SpeedtestCLIError(get_exception())
if args.list: if args.list:
try: try:
speedtest.get_servers() speedtest.get_servers()
except ServersRetrievalError: except (ServersRetrievalError, HTTP_ERRORS):
print_('Cannot retrieve speedtest server list') print_('Cannot retrieve speedtest server list')
raise raise SpeedtestCLIError(get_exception())
for _, servers in sorted(speedtest.servers.items()): for _, servers in sorted(speedtest.servers.items()):
for server in servers: for server in servers:
@ -1358,9 +1364,9 @@ def shell():
speedtest.get_servers(servers) speedtest.get_servers(servers)
except NoMatchedServers: except NoMatchedServers:
raise SpeedtestCLIError('No matched servers: %s' % args.server) raise SpeedtestCLIError('No matched servers: %s' % args.server)
except ServersRetrievalError: except (ServersRetrievalError, HTTP_ERRORS):
print_('Cannot retrieve speedtest server list') print_('Cannot retrieve speedtest server list')
raise raise SpeedtestCLIError(get_exception())
except InvalidServerIDType: except InvalidServerIDType:
raise SpeedtestCLIError('%s is an invalid server type, must ' raise SpeedtestCLIError('%s is an invalid server type, must '
'be an int' % args.server) 'be an int' % args.server)