diff --git a/stetho/agent/api.py b/stetho/agent/api.py index 92931b0..f7e38a3 100644 --- a/stetho/agent/api.py +++ b/stetho/agent/api.py @@ -16,6 +16,7 @@ import re from netaddr import IPNetwork from stetho.agent.common import utils as agent_utils +from stetho.agent.drivers import iperf as iperf_driver from stetho.agent.common import log LOG = log.get_logger() @@ -127,3 +128,35 @@ class AgentApi(object): # execute failed. message = stdout.pop(0) return agent_utils.make_response(code=stdcode, message=message) + + def setup_iperf_server(self, protocol='TCP', port=5001, window=None): + """iperf -s + """ + iperf = iperf_driver.IPerfDriver() + try: + data = iperf.start_server(protocol='TCP', port=5001, window=None) + return agent_utils.make_response(code=0, data=data) + except: + message = 'Start iperf server failed!' + return agent_utils.make_response(code=1, message=message) + + def teardown_iperf_server(self, pid): + iperf = iperf_driver.IPerfDriver() + try: + iperf.stop_server(pid) + return agent_utils.make_response(code=0) + except Exception as e: + message = e.message + return agent_utils.make_response(code=1, message=message) + + def start_iperf_client(self, host, protocol='TCP', timeout=5, + parallel=None, bandwidth=None): + iperf = iperf_driver.IPerfDriver() + try: + data = iperf.start_client(host, protocol='TCP', timeout=5, + parallel=None, bandwidth=None) + data['server_ip'] = host + return agent_utils.make_response(code=0, data=data) + except Exception as e: + message = e.message + return agent_utils.make_response(code=1, message=message) diff --git a/stetho/agent/common/utils.py b/stetho/agent/common/utils.py index 49c77ce..8afb26f 100644 --- a/stetho/agent/common/utils.py +++ b/stetho/agent/common/utils.py @@ -45,13 +45,20 @@ def execute(cmd, shell=False, root=False, timeout=10): def list_strip(lines): return [line.strip() for line in lines] - return stdcode, list_strip(stderr) if stdcode else list_strip(stdout) except Exception as e: LOG.error(e) raise +def execute_wait(cmd, shell=False, root=False): + subproc = subprocess.Popen(cmd, stdout=subprocess.PIPE, + stderr=subprocess.PIPE, shell=shell) + stdout, stderr = subproc.communicate() + stdcode = subproc.returncode + return stdcode, stdout, stderr + + def create_deamon(cmd, shell=False, root=False): """Usage: Create servcice process. @@ -62,7 +69,6 @@ def create_deamon(cmd, shell=False, root=False): LOG.info(cmd) subproc = subprocess.Popen(cmd, shell=shell, stdout=subprocess.PIPE, stderr=subprocess.PIPE) - stdcode = subproc.returncode return subproc.pid except Exception as e: LOG.error(e) @@ -70,7 +76,9 @@ def create_deamon(cmd, shell=False, root=False): def kill_process_by_id(pid): - os.kill(int(pid), signal.SIGKILL) + pid = int(pid) + os.kill(pid, signal.SIGILL) + os.waitpid(pid, 0) def get_interface(interface): diff --git a/stetho/agent/drivers/iperf.py b/stetho/agent/drivers/iperf.py index 51a8b83..6cacd5f 100644 --- a/stetho/agent/drivers/iperf.py +++ b/stetho/agent/drivers/iperf.py @@ -27,50 +27,41 @@ class IPerfDriver(object): def start_server(self, protocol='TCP', port=5001, mss=None, window=None): """iperf -s -D --mss mss """ - id = uuid.uuid4() - output = OUT_DIR + 'iperf-server-%s' % id - utils.replace_file(output) - cmd = ['iperf', '-s', '-D', '-p', str(port), '-o', output] + cmd = ['iperf', '-s', '-p', str(port)] if not cmp(protocol, 'UDP'): cmd.append('-u') if mss: - cmd.extend(['-m', str(mss)]) + cmd.extend(['-M', str(mss)]) if window: cmd.extend(['-w', str(window)]) - cmd.extend(['>', output]) pid = utils.create_deamon(cmd) data = dict() - data['id'] = id data['pid'] = pid return data def stop_server(self, pid): utils.kill_process_by_id(pid) - def start_client(self, host, port=5001, protocol='TCP', time=60, + def start_client(self, host, port=5001, protocol='TCP', timeout=5, parallel=None, bandwidth=None): """iperf -D -c host -t 60 """ - id = uuid.uuid4() - output = OUT_DIR + 'iperf-client-%s' % id - utils.replace_file(output) - cmd = ['iperf', '-D', '-c', host, '-p', str(port), '-t', str(time)] + cmd = ['iperf', '-c', host, '-p', str(port), '-t', str(timeout)] if not (protocol, 'UDP'): cmd.append('-u') if parallel: cmd.extend(['-P', str(parallel)]) if bandwidth: cmd.extend(['-b', '%sM' % bandwidth]) - cmd.extend(['>', output]) - utils.create_deamon(cmd) - data = dict() - data['id'] = id - return data - - def get_server_output(self, id): - # TODO: some analysis - pass - - def get_client_output(self, id): - # TODO: some analysis - pass + stdcode, stdout, stderr = utils.execute_wait(cmd) + if (not stdcode) or (not stderr): + out_dict = stdout.split('\n') + if not out_dict[-1]: + out_dict.pop() + out_data = out_dict[-1].split() + data = dict() + data['Bandwidth'] = out_data[-2] + out_data[-1] + data['Transfer'] = out_data[-4] + out_data[-3] + data['Interval'] = out_data[-6] + return data + raise Exception('Start iperf failed, please check on the node.') diff --git a/stetho/tests/unit/agent/common/test_utils.py b/stetho/tests/unit/agent/common/test_utils.py index 9c05a8b..fa4dc15 100644 --- a/stetho/tests/unit/agent/common/test_utils.py +++ b/stetho/tests/unit/agent/common/test_utils.py @@ -89,5 +89,6 @@ class TestUtils(unittest.TestCase): def test_kill_process_by_id(self): pid = 100 os.kill = mock.Mock() + os.waitpid = mock.Mock() utils.kill_process_by_id(pid) self.assertEqual(os.kill.called, True) diff --git a/stetho/tests/unit/agent/drivers/test_iperf.py b/stetho/tests/unit/agent/drivers/test_iperf.py index 4e2f15c..fce5bb8 100644 --- a/stetho/tests/unit/agent/drivers/test_iperf.py +++ b/stetho/tests/unit/agent/drivers/test_iperf.py @@ -29,6 +29,7 @@ class TestIPerfDriver(unittest.TestCase): self.assertEqual(data['pid'], 1000) def test_start_client(self): - utils.create_deamon = mock.Mock() + stdout = '[ 3] 0.0- 3.0 sec 497 MBytes 1.39 Gbits/sec' + utils.execute_wait = mock.Mock(return_value=(0, stdout, '')) self.iperfd.start_client('127.0.0.1') - self.assertEqual(utils.create_deamon.called, True) + self.assertEqual(utils.execute_wait.called, True) diff --git a/stetho/tests/unit/agent/test_api.py b/stetho/tests/unit/agent/test_api.py index 52e8451..c1feb14 100644 --- a/stetho/tests/unit/agent/test_api.py +++ b/stetho/tests/unit/agent/test_api.py @@ -17,6 +17,7 @@ import mock import unittest from stetho.agent import api from stetho.agent.common import utils as agent_utils +from stetho.agent.drivers import iperf as iperf_driver class TestApi(unittest.TestCase): @@ -67,3 +68,19 @@ class TestApi(unittest.TestCase): agent_utils.execute = mock.Mock(return_value=(1, stdout)) self.agent_api.teardown_link('eth0') self.assertEqual(agent_utils.make_response.called, True) + + def test_start_iperf_client(self): + agent_utils.create_deamon = mock.Mock(return_value=100) + self.agent_api.setup_iperf_server('UDP') + self.assertEqual(agent_utils.make_response.called, True) + + def test_teardown_iperf_server(self): + agent_utils.kill_process_by_id = mock.Mock() + self.agent_api.setup_iperf_server(100) + self.assertEqual(agent_utils.make_response.called, True) + + def test_start_client(self): + stdout = '[ 3] 0.0- 3.0 sec 497 MBytes 1.39 Gbits/sec' + agent_utils.execute_wait = mock.Mock(return_value=(0, stdout, '')) + self.agent_api.start_iperf_client(host='127.0.0.1') + self.assertEqual(agent_utils.make_response.called, True)