Add some UT test cases and fix pep8 errors for driver

JIRA:QTIP-89

Change-Id: I672ebaf733ec967307963a18f32ca0aa4afc672d
Signed-off-by: zhifeng.jiang <jiang.zhifeng@zte.com.cn>

2 files changed, 83 insertions, 20 deletions

diff --git a/func/driver.py b/func/driver.py
index 48c09c5d..33dbe320 100644
--- a/func/driver.py
+++ b/func/driver.py
@@ -6,13 +6,11 @@
 # which accompanies this distribution, and is available at
 # http://www.apache.org/licenses/LICENSE-2.0
 ##############################################################################
-
-
-
 import os
 import json
 from collections import defaultdict
 
+
 class Driver:
 
     def __init__(self):
@@ -21,8 +19,8 @@ class Driver:
         print os.environ['PWD']
         self.dic_json = defaultdict()
 
-    def drive_bench(self, benchmark, roles, benchmark_fname, benchmark_detail = None, pip_dict = None, proxy_info = None):
-        roles= sorted(roles)
+    def drive_bench(self, benchmark, roles, benchmark_fname, benchmark_detail=None, pip_dict=None, proxy_info=None):
+        roles = sorted(roles)
         pip_dict = sorted(pip_dict)
         result_dir = 'results'
         benchmark_name = benchmark + '.yaml'
@@ -33,30 +31,30 @@ class Driver:
         self.dic_json['workingdir'] = str(os.environ['PWD'])
         self.dic_json['fname'] = str(benchmark_fname)
         self.dic_json['username'] = str('root')
-        
+
         for key in proxy_info.keys():
             self.dic_json[key] = proxy_info[key]
-                    
+
         if os.environ['INSTALLER_TYPE'] == str('joid'):
-            self.dic_json['username']=str('ubuntu')
+            self.dic_json['username'] = str('ubuntu')
         if os.environ['INSTALLER_TYPE'] == str('apex'):
-            self.dic_json['username']=str('heat-admin')
-        for k,v in benchmark_detail:
-            self.dic_json[k]=v
+            self.dic_json['username'] = str('heat-admin')
+        for k, v in benchmark_detail:
+            self.dic_json[k] = v
         for k, v in roles:
-            self.dic_json['role']=k
-            index=1
-            if benchmark_detail != None:
+            self.dic_json['role'] = k
+            index = 1
+            if benchmark_detail is not None:
                 for values in v:
-                    if k ==  '1-server':
+                    if k == '1-server':
                         print values, 'saving IP'
-                        self.dic_json['ip'+str(index)]= str(values)
+                        self.dic_json['ip' + str(index)] = str(values)
                         if pip_dict[0][1][0]:
-                            self.dic_json['privateip'+str(index)] = pip_dict[0][1]
+                            self.dic_json['privateip' + str(index)] = pip_dict[0][1]
                         if not pip_dict[0][1][0]:
-                           self.dic_json['privateip'+str(index)] = 'NONE'
-                        index= index+1
+                            self.dic_json['privateip' + str(index)] = 'NONE'
+                        index = index + 1
             dic_json = json.dumps(dict(self.dic_json.items()))
             print dic_json
             run_play = 'ansible-playbook ./benchmarks/playbooks/{0}  --private-key=./data/QtipKey -i ./data/hosts --extra-vars \'{1}\''.format(benchmark_name, dic_json)
-            status = os.system(run_play)
+            os.system(run_play)
diff --git a/tests/driver_test.py b/tests/driver_test.py
new file mode 100644
index 00000000..39adc939
--- /dev/null
+++ b/tests/driver_test.py
@@ -0,0 +1,65 @@
+import pytest
+import mock
+import os
+import json
+from func.driver import Driver
+
+
+class TestClass:
+    @pytest.mark.parametrize("test_input, expected", [
+        (["iperf",
+          [('host', ['10.20.0.13', '10.20.0.15'])],
+          "iperf_bm.yaml",
+          [('duration', 20), ('protocol', 'tcp'), ('bandwidthGbps', 0)],
+          [("10.20.0.13", [None]), ("10.20.0.15", [None])],
+          {'http_proxy': 'http://10.20.0.1:8118',
+           'https_proxy': 'http://10.20.0.1:8118',
+           'no_proxy': 'localhost,127.0.0.1,10.20.*,192.168.*'},
+          'fuel'],
+         {'Dest_dir': 'results',
+          'ip1': '',
+          'ip2': '',
+          'installer': 'fuel',
+          'workingdir': '/home',
+          'fname': 'iperf_bm.yaml',
+          'username': 'root',
+          'http_proxy': 'http://10.20.0.1:8118',
+          'https_proxy': 'http://10.20.0.1:8118',
+          'no_proxy': 'localhost,127.0.0.1,10.20.*,192.168.*',
+          'duration': 20,
+          'protocol': 'tcp',
+          'bandwidthGbps': 0,
+          "role": "host"}),
+        (["iperf",
+          [('1-server', ['10.20.0.13']), ('2-host', ['10.20.0.15'])],
+          "iperf_vm.yaml",
+          [('duration', 20), ('protocol', 'tcp'), ('bandwidthGbps', 0)],
+          [("10.20.0.13", [None]), ("10.20.0.15", [None])],
+          {},
+          'joid'],
+         {'Dest_dir': 'results',
+          'ip1': '10.20.0.13',
+          'ip2': '',
+          'installer': 'joid',
+          "privateip1": "NONE",
+          'workingdir': '/home',
+          'fname': 'iperf_vm.yaml',
+          'username': 'ubuntu',
+          'duration': 20,
+          'protocol': 'tcp',
+          'bandwidthGbps': 0,
+          "role": "2-host"})
+    ])
+    @mock.patch('func.driver.os.system')
+    def test_driver_success(self, mock_system, test_input, expected):
+        mock_system.return_value = True
+        k = mock.patch.dict(os.environ, {'INSTALLER_TYPE': test_input[6], 'PWD': '/home'})
+        k.start()
+        dri = Driver()
+        dri.drive_bench(test_input[0], test_input[1], test_input[2], test_input[3], test_input[4], test_input[5])
+        call = mock_system.call_args
+        k.stop()
+        call_args, call_kwargs = call
+        real_call = call_args[0].split('extra-vars \'')[1]
+        real_call = real_call[0: len(real_call) - 1]
+        assert json.loads(real_call) == json.loads(json.dumps(expected))