cowmesh-network-test/cowmesh_pi_iperf_test.py

"""
similar to cowmesh_iperf_test but instead of running iperf router-to-router,
it runs iperf between computers which are connected to the network
(pi & nuc etc.)
"""
import datetime
import os
import re
import time
import asyncio, asyncssh, sys
import paramiko
import subprocess
import json
from cowmesh_helpers import cleanup_iperf_server

PROJECT_PATH = os.path.abspath(os.path.dirname(__file__))
SECRETS_PATH = os.path.join(PROJECT_PATH, "secrets.json")
with open(SECRETS_PATH, 'r') as f:
    SECRETS = json.loads(f.read())

base_node = "janastunuc"  # nuc in jaaga

nodes = [
    "janastunuc",   # nuc in jaaga
    "solipi",       # pi in guard
    "writer",       # pi in new-gazebo
]


class CowmeshPiIperfTester:

    def __init__(self, log=None, debug=False, time=10):
        if log:
            self.log = log
        self.debug = debug
        self.time = time
        self.results = {}

    async def log(self, msg):
        print(msg)

    async def debug_log(self, msg):
        if self.debug:
            await self.log(msg)

    async def test_between_two_nodes(self, node_a, node_b):
        await self.log("++ running test from {} to {}".format(node_a, node_b))
        user_a = SECRETS["HOST_INFO"][node_a]["user"]
        password_a = SECRETS["HOST_INFO"][node_a]["password"]

        myconn = paramiko.SSHClient()
        myconn.set_missing_host_key_policy(paramiko.AutoAddPolicy())

        ip_a = SECRETS["HOST_INFO"][node_a]["ip"]
        myconn.connect(ip_a, username=user_a, password=password_a)

        ip_b = SECRETS["HOST_INFO"][node_b]["ip"]

        remote_cmd = 'iperf -c {ip_b} -p 5001 -t {time}'.format(ip_b=ip_b, time=self.time)
        (stdin, stdout, stderr) = myconn.exec_command(remote_cmd)
        output = str(stdout.read())
        await self.debug_log("output: {}".format(output))
        await self.debug_log("errors: {}".format(stderr.read()))
        myconn.close()

        match = re.search("(\S+) Mbits", output)
        if match:
            to_return = match.group(1)
        else:
            to_return = None

        return to_return

    async def start_iperf_servers(self):
        for node in nodes:
            try:
                print("++ starting iperf server on {}".format(node))
                user = SECRETS["HOST_INFO"][node]["user"]
                password = SECRETS["HOST_INFO"][node]["password"]

                myconn = paramiko.SSHClient()
                myconn.set_missing_host_key_policy(paramiko.AutoAddPolicy())

                myconn.connect(node, username=user, password=password)

                remote_cmd = 'iperf -s &'
                (stdin, stdout, stderr) = myconn.exec_command(remote_cmd)
                await self.debug_log("{}".format(stdout.read()))
                await self.debug_log("{}".format(type(myconn)))
                await self.debug_log("Options available to deal with the connections are many like\n{}".format(dir(myconn)))
                myconn.close()
            except:
                try:
                    self.log("++ error starting iperf server on {}".format(node))
                except:
                    continue

    async def cleanup_iperf_servers(self):
        await self.log("++ shutting down iperf servers")
        for node in nodes:
            user = SECRETS["HOST_INFO"][node]["user"]
            password = SECRETS["HOST_INFO"][node]["password"]
            ip = SECRETS["HOST_INFO"][node]["ip"]
            await cleanup_iperf_server(node=node, ip=ip, username=user, password=password, log=self.log)

    async def run_test(self):

        await self.start_iperf_servers()

        for node_a in nodes:
            for node_b in nodes:
                if node_a == node_b:
                    await self.debug_log("skip self")
                    continue

                result_key = "{} -> {}".format(node_a, node_b)
                try:
                    r = await self.test_between_two_nodes(node_a, node_b)
                    self.results[result_key] = r
                except Exception as e:
                    self.results[result_key] = "error: {}".format(e)

        # try:
        #     await self.cleanup_iperf_servers()
        # except:
        #     await self.log("error shutting down iperf servers")

    async def output_results(self):
        results_str = ""
        now = datetime.datetime.now()
        date = now.date()
        time = now.time()
        results_str += "**** computer-to-computer iperf results on {date:%m-%d-%Y} at {time:%H:%M}:\n\n".format(date=date, time=time)
        for test_name, result in self.results.items():
            results_str += "{}: {} mbps\n".format(test_name, result)
        await self.log(results_str)


if __name__ == "__main__":
    try:
        tester = CowmeshPiIperfTester()

        async def main_fun():
            await tester.run_test()
            await tester.output_results()

        asyncio.get_event_loop().run_until_complete(main_fun())


    except (OSError, asyncssh.Error) as exc:
        sys.exit('SSH connection failed: ' + str(exc))