network_stress.py

#!/usr/bin/env python
import sys
import time
import psutil
import subprocess

wan_interface = subprocess.check_output([
    '/bin/bash',
    '-c',
    'ifconfig | grep $(host $(hostname) | cut -d \  -f 4) -B1 | '
    'head -n1 | cut -d \  -f 1']).strip()

get_counter = lambda: psutil.net_io_counters(True)[wan_interface].bytes_sent

def start_test(seconds):
    print 'Sending data...'

    # Record network counter and time before starting iperf
    start_bytes = get_counter()
    start_time = time.time()

    # Start iperf to increase traffic
    output = subprocess.check_output(['iperf', '-c', 'uk2.put.io',
                                      '-t', seconds, '-P', '10'])

    # Record network counter and time after iperf
    end_bytes = get_counter()
    end_time = time.time()

    # Getting the traffic created by iperf from the scripts output
    test_mbits_per_second = float(output.split('\n')[-2].split(' ')[-2])

    # Calculating total wan out traffic during the test
    bytes_sent = end_bytes - start_bytes
    bits_sent = bytes_sent * 8
    time_spent = end_time - start_time
    host_mbits_per_second = (bits_sent / time_spent) / (1000 * 1000)

    print 'Test completed. Speed average %s Mbits/s (from test: %s)' % (
        host_mbits_per_second, test_mbits_per_second)


if __name__ == '__main__':
    start_test(sys.argv[1])