Проверка скорости интернета библиотекой Requests в мультипроцессинге

Доброго времени, уважаемые жители Хабра!

Сегодня речь пойдет о том, как из идеи по замеру скорости, был создан скрипт для загрузки файла изображения и отправки его же обратно на сервер, с расчетом времени выполнения каждой из функций и вычисления скорости.

Начну со списка используемых библиотек:

• import os
• from multiprocessing import Pool
• import time
• import pandas as pd
• import requests

Далее нам нужен список серверов, я предпочел создать для этого словарь:

server_list = [
    {
        'server_id': 3682,
        'download': 'http://moscow.speedtest.rt.ru:8080/speedtest/random7000x7000.jpg',
        'upload': 'http://moscow.speedtest.rt.ru:8080/speedtest/upload.php'
    }
]

Пишем первую функцию:

def download(id, path):
    start = time.time()
    file_name = str(id) + str(path.split('/')[-1])
    try:
        r = requests.get(path, stream=True, timeout=5)
    except:
        return 0
    size = int(r.headers.get('Content-Length', 0))
    with open(file_name, 'wb') as f:
        for chunk in r.iter_content(chunk_size=1024):
            if chunk:
                f.write(chunk)

    end = time.time()
    duration = end - start
    sp = (((size * 8) / 1024) / 1024) / duration

    return sp

Теперь подробнее о том, что происходит.

В функции есть время старта и время окончания(в секундах), из которых в дальнейшем мы получаем время жизни. В имя файла записываем id сервера и название изображения(сделано для того, чтобы не возникало конфликтов при загрузке из множества источников). Далее делаем GET запрос, получаем размер файла(в байтах) и сохраняем его на диск. Переводим байты в биты, еще немного магии с формулами и на выходе имеем скорость в MBit/s.

Следующая функция — отдача файла на сервер:

def upload(id, path):
    start = time.time()
    file_name = str(id) + 'random7000x7000.jpg'
    with open(file_name, 'rb') as f:
        files = {'Upload': (file_name, f.read())}
    try:
        requests.post(path, files=files)
    except:
        return 0
    size = os.path.getsize(file_name)
    end = time.time()
    duration = end - start
    sp = (((size * 8) / 1024) / 1024) / duration

    return sp

Здесь принцип тот же, только мы берем файл из локальной папки и POST запросом отправляем.

Наша следующая задача получить данные из двух предыдущих функций. Пишем еще одну функцию:

def test_f(conn, server):
    speed_download = download(server['server_id'], server['download'])
    speed_upload = upload(server['server_id'], server['upload'])
    return server['server_id'], speed_download, speed_upload
    

Осталось дело за малым, прикрутить мультипроцессинг с пулом и параллельной функцией map:

def main():
    pool = Pool()
    data = pool.map(test_f, server_list)

    df = pd.DataFrame(data, columns=['Server', 'Download', 'Upload'])
    print(df)

    pool.close()
    pool.join()

if __name__ == '__main__':
    main()

Скрипт готов к использованию, для удобства вывода я использовал библиотеку pandas. Так же вывод можно поместить в базу и собирать статистику для анализа.

Спасибо за внимание!

UPD: Поправил исключения, внес правки в работу мультипроцессинга(заменил цикл на параллельную функцию), добавил таймаут для GET запроса