speedtest.py

# -*-: coding: utf-8 -*-

import asyncio
import aiohttp
import traceback
import datetime
import math

def hsize(s):
    suffixs = ["B", "KB", "MB", "GB"]
    size = s
    for suf in suffixs:
        if size < 1024:
            break
        size = size/1024
    return "{}{}".format(math.ceil(size), suf)

def htime(t):
    suffixs = ["sec", "min", "hour"]
    time = t
    for suf in suffixs:
        if time < 60:
            break
        time = time/60
    return "{}{}".format(math.ceil(time), suf)

async def meaured_download(did, session, url):
    print("{} trying to start download...".format(did))
    download_size = 0
    speed = 0
    start_ts = datetime.datetime.now().timestamp()
    try:
        async with session.get(url) as resp:
            if resp.status >= 200 and resp.status < 300:
                while True:
                    data = await resp.content.read(2000)
                    data_len = len(data)
                    if data_len <= 0:
                        break
                    download_size = download_size + data_len
    except Exception as err:
        print("{} got error while downloading...".format(did))
        traceback.print_exc()
    end_ts = datetime.datetime.now().timestamp()
    total_time = end_ts - start_ts
    speed = int(download_size / total_time)
    print("{} finished, total time is {}, total download size is {}, average speed is {}/s".format(did, htime(total_time), hsize(download_size), hsize(speed)))
    return (download_size, speed)

def main(wid, url, conn_count):
    print("woker {} started".format(wid))
    tasks = []
    loop_count = int(conn_count)
    loop = asyncio.get_event_loop()
    session = aiohttp.ClientSession(connector=aiohttp.TCPConnector(limit=0), read_timeout=0)
    for i in range(loop_count):
        tasks.append(asyncio.ensure_future(meaured_download(i, session, url)))
    loop.run_until_complete(asyncio.gather(*tasks))
    print("worker {} finished".format(wid))

def usage(pname):
    print("Usage: {} url conn_count worker_count".format(pname))

if __name__ == "__main__":
    import sys
    if len(sys.argv) != 4:
        usage(sys.argv[0])
        exit(1)

    worker_count = int(sys.argv[3])
    from multiprocessing import Process
    p = []
    for i in range(worker_count):
        proc = Process(target=main, args=(i, sys.argv[1], sys.argv[2]))
        proc.start()
        p.append(proc)

    for i in range(worker_count):
        p[i].join()