refactor code related with setup env

- modify env.py and scripts
- add unit test

Change-Id: I416cd517bdddfeb907675705009d83ecada7f87c
Signed-off-by: wu.zhihui <wu.zhihui1@zte.com.cn>

4 files changed, 175 insertions, 112 deletions

diff --git a/qtip/scripts/fetch_compute_ips.sh b/qtip/scripts/generate_host_file.sh
index a2618ec9..1353cacd 100755
--- a/qtip/scripts/fetch_compute_ips.sh
+++ b/qtip/scripts/generate_host_file.sh
@@ -10,17 +10,17 @@
 
 
 usage(){
-   echo "usage: $0 [-v] -i <installer_type> -a <installer_ip>" >&2
+   echo "usage: $0 [-v] -i <installer_type> -a <installer_ip> -d <host_file>" >&2
    echo "[-v] Virtualized deployment" >&2
 }
 
 info()  {
-   logger -s -t "fetch_compute_info.info" "$*"
+   logger -s -t "generate_host_file.info" "$*"
 }
 
 
 error() {
-   logger -s -t "fetch_compute_info.error" "$*"
+   logger -s -t "generate_host_file.error" "$*"
    exit 1
 }
 
@@ -44,6 +44,7 @@ while getopts ":i:a:h:v" optchar; do
    case "${optchar}" in
        i) installer_type=${OPTARG} ;;
        a) installer_ip=${OPTARG} ;;
+       d) host_file=${OPTARG} ;;
        v) DEPLOY_TYPE="virt" ;;
        *) echo "Non-option argument: '-${OPTARG}'" >&2
           usage
@@ -108,10 +109,11 @@ if [ -z "$IPS" ]; then
    error "The compute node $IPS are not up. Please check that the POD is correctly deployed."
 else
    echo "-------- all compute node ips: --------"
-   rm $HOME/ips.log
-   touch $HOME/ips.log
-   echo "$IPS" > $HOME/qtip/ips.log
-   echo $IPS
+   rm $host_file
+   touch $host_file
+   echo "[hosts]" >> $host_file
+   echo "$IPS" >> $host_file
+   cat $host_file
 fi
 
 exit 0
diff --git a/qtip/scripts/qtip_creds.sh b/qtip/scripts/qtip_creds.sh
index af051ac5..d338115f 100755
--- a/qtip/scripts/qtip_creds.sh
+++ b/qtip/scripts/qtip_creds.sh
@@ -1,30 +1,34 @@
 #! /bin/bash
 
 DEST_IP=$1
+PRIVATE_KEY=$2
+PUBLIC_KEY=$2.pub
+KEYNAME=$(basename PRIVATE_KEY)
+
 echo $INSTALLER_TYPE
 echo $INSTALLER_IP
 sshoptions="-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null"
 case "$INSTALLER_TYPE" in
     apex)
-        scp $sshoptions -i $APEX_KEY ./config/QtipKey.pub stack@$INSTALLER_IP:/home/stack
-        scp $sshoptions -i $APEX_KEY ./config/QtipKey stack@$INSTALLER_IP:/home/stack
-        ssh $sshoptions -i $APEX_KEY stack@$INSTALLER_IP "ssh-copy-id $sshoptions -i /home/stack/QtipKey.pub heat-admin@$DEST_IP && rm -rf /home/stack/QtipKey && rm -rf /home/stack/QtipKey.pub"
+        scp $sshoptions -i $APEX_KEY $PUBLIC_KEY stack@$INSTALLER_IP:/home/stack
+        scp $sshoptions -i $APEX_KEY $PRIVATE_KEY stack@$INSTALLER_IP:/home/stack
+        ssh $sshoptions -i $APEX_KEY stack@$INSTALLER_IP "ssh-copy-id $sshoptions -i /home/stack/$KEYNAME.pub heat-admin@$DEST_IP && rm -rf /home/stack/$KEYNAME && rm -rf /home/stack/$KEYNAME.pub"
         ;;
     fuel)
         PSWD="r00tme"
-        sshpass -p $PSWD scp $sshoptions ./config/QtipKey.pub root@$INSTALLER_IP:/root
-        sshpass -p $PSWD scp $sshoptions ./config/QtipKey root@$INSTALLER_IP:/root
+        sshpass -p $PSWD scp $sshoptions $PUBLIC_KEY root@$INSTALLER_IP:/root
+        sshpass -p $PSWD scp $sshoptions $PRIVATE_KEY root@$INSTALLER_IP:/root
         sshpass -p $PSWD ssh $sshoptions root@$INSTALLER_IP "grep -q '\-F /dev/null ' /usr/bin/ssh-copy-id || sed -i 's/\(ssh -i.*$\)/\1\n -F \/dev\/null \\\/g' `which ssh-copy-id`"
-        sshpass -p $PSWD ssh $sshoptions root@$INSTALLER_IP "ssh-copy-id $sshoptions -i /root/QtipKey root@$DEST_IP && rm -rf /root/QtipKey && rm -rf /root/QtipKey.pub"
+        sshpass -p $PSWD ssh $sshoptions root@$INSTALLER_IP "ssh-copy-id $sshoptions -i /root/$KEYNAME root@$DEST_IP && rm -rf /root/$KEYNAME && rm -rf /root/$KEYNAME.pub"
         ;;
     compass)
         PSWD="root"
-        sshpass -p $PSWD scp $sshoptions ./config/QtipKey.pub root@$INSTALLER_IP:/root
-        sshpass -p $PSWD scp $sshoptions ./config/QtipKey root@$INSTALLER_IP:/root
-        sshpass -p $PSWD ssh $sshoptions root@$INSTALLER_IP "ssh-copy-id $sshoptions -i /root/QtipKey.pub root@$DEST_IP && rm -rf /root/QtipKey && rm -rf /root/QtipKey.pub"
+        sshpass -p $PSWD scp $sshoptions $PUBLIC_KEY root@$INSTALLER_IP:/root
+        sshpass -p $PSWD scp $sshoptions $PRIVATE_KEY root@$INSTALLER_IP:/root
+        sshpass -p $PSWD ssh $sshoptions root@$INSTALLER_IP "ssh-copy-id $sshoptions -i /root/$KEYNAME.pub root@$DEST_IP && rm -rf /root/$KEYNAME && rm -rf /root/$KEYNAME.pub"
         ;;
     joid)
         PSWD="joid";;
        *)
         echo "Unkown installer $INSTALLER_TYPE specified";;
-esac
+esac
\ No newline at end of file
diff --git a/qtip/util/env.py b/qtip/util/env.py
index 4e7a31c7..0585a4c1 100644
--- a/qtip/util/env.py
+++ b/qtip/util/env.py
@@ -7,131 +7,94 @@
 # http://www.apache.org/licenses/LICENSE-2.0
 ##############################################################################
 import os
-import paramiko
 import socket
 import time
 from os import path
-from os.path import expanduser
+
+import paramiko
 
 SCRIPT_DIR = path.join(path.dirname(__file__), path.pardir, 'scripts')
-CONFIG_DIR = path.join(path.dirname(__file__), path.pardir, path.pardir,
-                       'config')
-PRIVATE_KEY = CONFIG_DIR + '/QtipKey'
-PUBLIC_KEY = CONFIG_DIR + '/QtipKey.pub'
-IPS_FILE = expanduser('~') + "/qtip/ips.log"
-HOST_FILE = CONFIG_DIR + "/host"
+KEYNAME = 'QtipKey'
+PRIVATE_KEY = '{0}/qtip/{1}'.format(os.environ['HOME'], KEYNAME)
+PUBLIC_KEY = PRIVATE_KEY + '.pub'
+HOST_FILE = '{0}/qtip/hosts'.format(os.environ['HOME'])
 
 
-def fetch_compute_ips_via_installer():
-    clean_file(IPS_FILE)
+def all_files_exist(*files):
+    if len(files) == 0:
+        return False
+    flag = True
+    for f_item in files:
+        flag &= path.isfile(f_item)
+        print("Is {0} existed: {1}".format(f_item, flag))
+    return flag
 
+
+def clean_file(*files):
+    if len(files) == 0:
+        print('Nothing to clean')
+        return False
+
+    def clean(f):
+        try:
+            if all_files_exist(f):
+                os.remove(f)
+                print("Removed: {0}".format(f))
+            else:
+                print("Not exists: {0}".format(f))
+            return True
+        except OSError as error:
+            print("Not able to Remove: {0}".format(f), error)
+            return False
+
+    results = map(clean, files)
+    return len(results) == len(files) and False not in results
+
+
+def generate_host_file(hostfile=HOST_FILE):
     installer_type = str(os.environ['INSTALLER_TYPE'].lower())
     installer_ip = str(os.environ['INSTALLER_IP'])
+
     if installer_type not in ["fuel"]:
-        raise RuntimeError("%s is not supported" % installer_type)
+        raise ValueError("%s is not supported" % installer_type)
     if not installer_ip:
-        raise RuntimeError("undefine environment variable INSTALLER_IP")
+        raise ValueError("The value of environment variable INSTALLER_IP is empty")
 
-    cmd = "bash %s/fetch_compute_ips.sh -i %s -a %s" % \
-        (SCRIPT_DIR, installer_type, installer_ip)
+    cmd = "bash %s/generate_host_file.sh -i %s -a %s -d %s" % \
+        (SCRIPT_DIR, installer_type, installer_ip, hostfile)
     os.system(cmd)
-    if path.isfile(IPS_FILE):
-        return True
-    else:
-        return False
+    return all_files_exist(hostfile)
 
 
-def parse_ips():
-    ip_list = []
-    with open(IPS_FILE, "r") as outfile:
-        data = outfile.read()
-        if data:
-            ip_list.extend(data.rstrip('\n').split('\n'))
-    return ip_list
+def generate_keypair(keyname='QtipKey'):
+    """Generating ssh keypair"""
+    cmd = "ssh-keygen -t rsa -N "" -f {0} -q -b 2048".format(keyname)
+    os.system(cmd)
+    return all_files_exist(PRIVATE_KEY, PUBLIC_KEY)
 
 
-def ssh_test(ip):
+def pass_keypair(ip, private_key=PRIVATE_KEY):
     os.system('ssh-keyscan %s >> /root/.ssh/known_hosts' % ip)
     time.sleep(2)
 
-    ssh_cmd = '%s/qtip_creds.sh %s' % (SCRIPT_DIR, ip)
+    ssh_cmd = '%s/qtip_creds.sh %s %s' % (SCRIPT_DIR, ip, private_key)
     os.system(ssh_cmd)
 
+
+def ssh_is_ok(ip, private_key=PRIVATE_KEY, attempts=100):
     ssh = paramiko.SSHClient()
     ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
-    ssh.connect(ip, key_filename='{0}/QtipKey'.format(CONFIG_DIR))
+    ssh.connect(ip, key_filename=private_key)
 
-    for attempts in range(100):
+    for attempt in range(attempts):
         try:
             stdin, stdout, stderr = ssh.exec_command('uname')
             if not stderr.readlines():
-                print("{0}: SSH test successful")
+                print("{0}: SSH test successful".format(ip))
                 return True
         except socket.error:
-            if attempts == 99:
+            if attempt == (attempts - 1):
                 return False
+            print("%s times ssh test......failed" % attempt)
             time.sleep(2)
-
-
-def ping_test(ip, attempts=30):
-    ping_cmd = 'ping -D -c1 {0}'.format(ip)
-    for i in range(attempts):
-        if os.system(ping_cmd):
-            print('\nWaiting for machine\n')
-            time.sleep(10)
-        else:
-            print('\n\n %s is UP \n\n ' % ip)
-            return True
-        if i == 29:
-            return False
-
-
-def check_nodes_connectivity():
-    ip_list = parse_ips()
-    for ip in ip_list:
-        if not ping_test(ip):
-            raise RuntimeError("{0}: Ping test failed".format(ip))
-        if not ssh_test(ip):
-            raise RuntimeError("{0}: SSH test failed".format(ip))
-
-
-def generate_host_file():
-    ip_list = parse_ips()
-    with open(HOST_FILE, 'w') as host_file:
-        for index, item in enumerate(ip_list):
-            host_file.write("[host_{0}]\n".format(index))
-            host_file.write(item + '\n')
-
-
-def generate_keypair():
-    """Generating ssh keypair"""
-    if not clean_keypair():
-        raise RuntimeError("Cann't remove old keypair")
-
-    cmd = "ssh-keygen -t rsa -N "" -f {0} -q".format(PRIVATE_KEY)
-    os.system(cmd)
-
-    if path.isfile(PRIVATE_KEY) and path.isfile(PUBLIC_KEY):
-        return True
-    else:
-        return False
-
-
-def clean_file(file_path):
-    try:
-        if path.isfile(file_path):
-            os.remove(file_path)
-            print("Removed: " + file_path)
-        else:
-            print("Not exists: " + file_path)
-    except OSError, error:
-        print("Not able to Remove: " + file_path, error)
-        return False
-    return True
-
-
-def clean_keypair():
-    flag = True
-    flag &= clean_file(PRIVATE_KEY)
-    flag &= clean_file(PUBLIC_KEY)
-    return flag
+    return False
diff --git a/tests/unit/util/env_test.py b/tests/unit/util/env_test.py
new file mode 100644
index 00000000..38ac988b
--- /dev/null
+++ b/tests/unit/util/env_test.py
@@ -0,0 +1,94 @@
+###############################################################
+# Copyright (c) 2017 ZTE Corporation
+#
+# All rights reserved. This program and the accompanying materials
+# are made available under the terms of the Apache License, Version 2.0
+# which accompanies this distribution, and is available at
+# http://www.apache.org/licenses/LICENSE-2.0
+##############################################################################
+
+import time
+
+import mock
+import pytest
+
+from qtip.util import env
+
+
+def test_all_files_exist(tmpdir):
+    exist_file = tmpdir.mkdir('qtip').join('hello.txt')
+    exist_file.write("hello")
+    non_exist_file = tmpdir.strpath + '/tmp.txt'
+    assert env.all_files_exist() is False
+    assert env.all_files_exist(str(exist_file))
+    assert env.all_files_exist(non_exist_file) is False
+    assert env.all_files_exist(str(exist_file), non_exist_file) is False
+
+
+def test_clean_file(tmpdir):
+    exist_file = tmpdir.mkdir('qtip').join('hello.txt')
+    exist_file.write("hello")
+    non_exist_file = tmpdir.strpath + '/tmp.txt'
+
+    assert env.clean_file() is False
+    assert env.clean_file(str(exist_file))
+    assert env.clean_file(non_exist_file)
+
+
+def test_generate_host_file_without_setenv(monkeypatch):
+    def setenv(*args):
+        monkeypatch.setenv('INSTALLER_TYPE', args[0])
+        monkeypatch.setenv('INSTALLER_IP', args[1])
+
+    with pytest.raises(KeyError) as excinfo:
+        env.generate_host_file()
+    assert 'INSTALLER_TYPE' in str(excinfo.value)
+
+    with pytest.raises(ValueError) as excinfo:
+        setenv('fuel_1', '10.20.0.2')
+        env.generate_host_file()
+    assert 'fuel_1 is not supported' in str(excinfo.value)
+
+    with pytest.raises(ValueError) as excinfo:
+        setenv('fuel', '')
+        env.generate_host_file()
+    assert 'The value of environment variable INSTALLER_IP is empty' \
+           in str(excinfo.value)
+
+
+def test_generate_host_file(monkeypatch, tmpdir):
+    monkeypatch.setenv('INSTALLER_TYPE', 'fuel')
+    monkeypatch.setenv('INSTALLER_IP', '10.20.0.2')
+    hostfile = tmpdir.mkdir('qtip').join('hosts')
+    hostfile.write('')
+    assert env.generate_host_file(str(hostfile))
+
+
+def test_generate_keypair():
+    with mock.patch('os.system') as mock_os:
+        env.generate_keypair()
+        assert mock_os.call_count == 1
+
+
+def test_pass_keypair(monkeypatch):
+    monkeypatch.setattr(time, 'sleep', lambda s: None)
+    with mock.patch('os.system') as mock_os:
+        env.pass_keypair('10.20.0.10')
+        assert mock_os.call_count == 2
+
+
+@pytest.mark.parametrize("stderrinfo, expected", [
+    ('', True),
+    ('sorry', False)
+])
+@mock.patch('paramiko.SSHClient')
+def test_ssh_is_ok(mock_sshclient, stderrinfo, expected):
+    stderr = mock.MagicMock()
+    stderr.readlines.return_value = stderrinfo
+    test_ssh_client = mock_sshclient.return_value
+    test_ssh_client.exec_command.return_value = ('', '', stderr)
+    result = env.ssh_is_ok('10.20.0.3')
+    assert result == expected
+    test_ssh_client.connect.assert_called_once_with(
+        '10.20.0.3', key_filename=env.PRIVATE_KEY)
+    test_ssh_client.exec_command.assert_called_with('uname')