Merge pull request #18 from changzhi1990/stethoclient

Upload iperf CLI

stetho/agent/api.py

@@ -150,7 +150,7 @@ class AgentApi(object):
             return agent_utils.make_response(code=1, message=message)
 
     def start_iperf_client(self, host, protocol='TCP', timeout=5,
-                           parallel=None, bandwidth=None):
+                           parallel=None, bandwidth=None, port=5001):
         iperf = iperf_driver.IPerfDriver()
         try:
             data = iperf.start_client(host, protocol='TCP', timeout=5,

stetho/agent/drivers/iperf.py

@@ -60,8 +60,8 @@ class IPerfDriver(object):
                 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['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.')

stetho/stethoclient/agent_api.py

@@ -21,7 +21,6 @@ import sys
 from cliff.command import Command
 from cliff.lister import Lister
 from json import JSONDecoder
-from stetho.stethoclient.constants import AGENT_INFOS
 
 LISTEN_PORT = 9698
 SETUP_LINK_IP_PRE = "192.168.100."
@@ -48,6 +47,16 @@ class Logger():
         print Logger.FAIL + info + Logger.ENDC
 
 
+try:
+    from stetho.stethoclient.constants import AGENT_INFOS
+except:
+    AGENT_INFOS = {
+        'agent-64': "127.0.0.1",
+        'agent-65': "127.0.0.1",
+    }
+    Logger.log_fail("Import stetho configure file fail. Use fake data!")
+
+
 def setup_server(agent):
     log = logging.getLogger(__name__)
     if agent in AGENT_INFOS:

stetho/stethoclient/drivers/iperf_api.py

@@ -0,0 +1,133 @@
+#!/usr/bin/python
+# Copyright 2015 UnitedStack, Inc.
+# All Rights Reserved.
+#
+#    Licensed under the Apache License, Version 2.0 (the "License"); you may
+#    not use this file except in compliance with the License. You may obtain
+#    a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+#    License for the specific language governing permissions and limitations
+#    under the License.
+
+import logging
+import jsonrpclib
+import socket
+import sys
+
+from cliff.lister import Lister
+
+LISTEN_PORT = 9698
+
+
+class Logger():
+    HEADER = '\033[95m'
+    OKBLUE = '\033[94m'
+    OKGREEN = '\033[92m'
+    WARNING = '\033[93m'
+    FAIL = '\033[91m'
+    ENDC = '\033[0m'
+
+    @staticmethod
+    def log_normal(info):
+        print Logger.OKBLUE + info + Logger.ENDC
+
+    @staticmethod
+    def log_high(info):
+        print Logger.OKGREEN + info + Logger.ENDC
+
+    @staticmethod
+    def log_fail(info):
+        print Logger.FAIL + info + Logger.ENDC
+
+try:
+    from stetho.stethoclient.constants import AGENT_INFOS
+except:
+    AGENT_INFOS = {
+        'agent-64': "127.0.0.1",
+        'agent-65': "127.0.0.1",
+    }
+    Logger.log_fail("Import stetho configure file fail. Use fake data!")
+
+
+def setup_server(agent):
+    log = logging.getLogger(__name__)
+    if agent in AGENT_INFOS:
+        log.debug('get agent:%s ip_address:%s' % (
+            agent, AGENT_INFOS[agent]))
+    else:
+        log.error('Agent %s not configured. Please check it.' % (
+                         agent))
+        sys.exit()
+    log.debug('Begin create connection with http://%s:9698.' % (
+                    agent))
+    server = jsonrpclib.Server('http://%s:%s' % (
+             AGENT_INFOS[agent], LISTEN_PORT))
+    log.debug('Create connection with %s success.' % (agent))
+    return server
+
+
+def get_ip_by_hostname(hostname):
+    # TODO(changzhi): DNS resolve
+    return socket.gethostbyname(hostname)
+
+
+class CheckIperf(Lister):
+    "Setup Iperf server in agent"
+
+    log = logging.getLogger(__name__)
+
+    def get_parser(self, prog_name):
+        parser = super(CheckIperf, self).get_parser(prog_name)
+        parser.add_argument('server_agent', default='bad')
+        parser.add_argument('client_agent', default='bad')
+        parser.add_argument('--server_protocol', nargs='?', default='TCP')
+        parser.add_argument('--server_port', nargs='?', default='5001')
+        parser.add_argument('--client_protocol', nargs='?', default='TCP')
+        parser.add_argument('--client_port', nargs='?', default='5001')
+        parser.add_argument('--client_timeout', nargs='?', default='5')
+        parser.add_argument('--client_parallel', nargs='?', default=None)
+        parser.add_argument('--client_bandwidth', nargs='?', default=None)
+        return parser
+
+    def take_action(self, parsed_args):
+        self.log.debug('Get parsed_args: %s' % parsed_args)
+        server = setup_server(parsed_args.server_agent)
+        client = setup_server(parsed_args.client_agent)
+        iperf_server_pdid = None
+        # setup iperf server
+        res = server.setup_iperf_server(protocol=parsed_args.server_protocol,
+                                        port=parsed_args.server_port)
+        self.log.debug('Response is %s' % res)
+        if res['code'] == 1:
+            Logger.log_fail(res['message'])
+            sys.exit()
+        if res['code'] == 0:
+            msg = (('Iperf server setup success and runs in '
+                    'pid:%s') % (res['data']['pid']))
+            Logger.log_high(msg)
+            iperf_server_pdid = res['data']['pid']
+        # setup iperf client
+        host = get_ip_by_hostname(parsed_args.server_agent)
+        res = client.start_iperf_client(protocol=parsed_args.client_protocol,
+                                        host=host,
+                                        timeout=parsed_args.client_timeout,
+                                        parallel=parsed_args.client_parallel,
+                                        bandwidth=parsed_args.client_bandwidth,
+                                        port=parsed_args.client_port)
+        self.log.debug('Response is %s' % res)
+        # kill iperf server
+        r = server.teardown_iperf_server(iperf_server_pdid)
+        if r['code'] == 1:
+            Logger.log_fail(r['message'])
+            sys.exit()
+        if r['code'] == 0:
+            msg = (('Iperf server delete success and '
+                    'pid:%s') % (iperf_server_pdid))
+            Logger.log_high(msg)
+        return (('Field', 'Value'),
+                ((k, v) for k, v in res['data'].items()))

stetho/stethoclient/shell.py

@@ -22,6 +22,7 @@ import logging
 from cliff import app
 from cliff import commandmanager
 from stetho.stethoclient import agent_api
+from stetho.stethoclient.drivers import iperf_api
 from stetho.stethoclient import strutils
 
 
@@ -36,6 +37,7 @@ COMMAND_V1 = {
     'check-ports-on-br': agent_api.CheckPortsOnBr,
     'get-interface': agent_api.GetInterface,
     'check-vlan-interface': agent_api.CheckVlanInterface,
+    'check-iperf': iperf_api.CheckIperf,
 }
 
 COMMANDS = {'0.1': COMMAND_V1}

stetho/tests/unit/stethoclient/test_stethoclient.py

@@ -2,6 +2,7 @@ import mock
 import unittest
 from stetho.stethoclient import shell
 from stetho.stethoclient import agent_api
+from stetho.stethoclient.drivers import iperf_api
 
 
 class TestStethoClientMethods(unittest.TestCase):
@@ -35,3 +36,9 @@ class TestStethoClientMethods(unittest.TestCase):
         agent_api.CheckPortsOnBr.take_action = mock.Mock(return_value=r)
         shell.main(['check-ports-on-br', 'agent-64', 'br0', 'a'])
         self.assertEqual(agent_api.CheckPortsOnBr.take_action.called, True)
+
+    def test_check_iperf(self):
+        r = (('Field', 'Value'), [])
+        iperf_api.CheckIperf.take_action = mock.Mock(return_value=r)
+        shell.main(['check-iperf', 'agent-64', 'agent-64'])
+        self.assertEqual(iperf_api.CheckIperf.take_action.called, True)