Merge "rework env.py"

3 files changed, 396 insertions, 79 deletions

diff --git a/qtip/util/env.py b/qtip/util/env.py
index 0585a4c1..24e08658 100644
--- a/qtip/util/env.py
+++ b/qtip/util/env.py
@@ -1,15 +1,18 @@
 ##############################################################################
-# Copyright (c) 2016 Dell Inc, ZTE and others.
+# Copyright (c) 2017 ZTE and others.
 #
 # 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
 ##############################################################################
+from collections import defaultdict
 import os
+from os import path
+import re
 import socket
+import sys
 import time
-from os import path
 
 import paramiko
 
@@ -51,50 +54,134 @@ def clean_file(*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 ValueError("%s is not supported" % installer_type)
-    if not installer_ip:
-        raise ValueError("The value of environment variable INSTALLER_IP is empty")
-
-    cmd = "bash %s/generate_host_file.sh -i %s -a %s -d %s" % \
-        (SCRIPT_DIR, installer_type, installer_ip, hostfile)
-    os.system(cmd)
-    return all_files_exist(hostfile)
-
-
-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 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 %s' % (SCRIPT_DIR, ip, private_key)
-    os.system(ssh_cmd)
+class AnsibleEnvSetup(object):
+    def __init__(self):
+        self.keypair = defaultdict(str)
+        self.hostfile = None
+        self.host_ip_list = []
 
+    def setup(self, config={}):
+        try:
+            if 'hostfile' in config:
+                self.check_hostfile(config['hostfile'])
+            else:
+                self.generate_default_hostfile()
+            self.fetch_host_ip_from_hostfile()
+            if 'keypair' in config:
+                self.check_keypair(config['keypair'])
+            else:
+                self.generate_default_keypair()
+            self.pass_keypair_to_remote()
+            self.check_hosts_ssh_connectivity()
+        except Exception as error:
+            print(error)
+            sys.exit(1)
+
+    def check_keypair(self, keypair):
+        self.keypair = defaultdict(str)
+        if all_files_exist(keypair, '{0}.pub'.format(keypair)):
+            self.keypair['private'] = keypair
+            self.keypair['public'] = '{0}.pub'.format(keypair)
+        else:
+            raise RuntimeError("The keypairs you in the configuration file"
+                               " is invalid or not existed.")
+
+    def generate_default_keypair(self):
+        if not all_files_exist(PRIVATE_KEY, PUBLIC_KEY):
+            print("Generate default keypair {0} under "
+                  "{1}".format(KEYNAME, os.environ['HOME']))
+            cmd = '''ssh-keygen -t rsa -N "" -f {0} -q -b 2048'''.format(
+                PRIVATE_KEY)
+            os.system(cmd)
+        self.keypair['private'] = PRIVATE_KEY
+        self.keypair['public'] = PUBLIC_KEY
+
+    def pass_keypair_to_remote(self):
+        results = map(lambda ip: self._pass_keypair(ip, self.keypair['private']),
+                      self.host_ip_list)
+
+        if not (len(results) == len(self.host_ip_list) and False not in results):
+            raise RuntimeError("Failed on passing keypair to remote.")
+
+    @staticmethod
+    def _pass_keypair(ip, private_key):
+        try:
+            os.system('ssh-keyscan %s >> /root/.ssh/known_hosts' % ip)
+            time.sleep(2)
+            ssh_cmd = '%s/qtip_creds.sh %s %s' % (SCRIPT_DIR, ip, private_key)
+            os.system(ssh_cmd)
+            print('Pass keypair to remote hosts {0} successfully'.format(ip))
+            return True
+        except Exception as error:
+            print(error)
+            return False
 
-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=private_key)
+    def check_hostfile(self, hostfile):
+        if all_files_exist(hostfile):
+            self.hostfile = hostfile
+        else:
+            raise RuntimeError(
+                "The hostfile {0} is invalid or not existed.".format(hostfile))
 
-    for attempt in range(attempts):
+    def generate_default_hostfile(self):
         try:
-            stdin, stdout, stderr = ssh.exec_command('uname')
-            if not stderr.readlines():
-                print("{0}: SSH test successful".format(ip))
-                return True
-        except socket.error:
-            if attempt == (attempts - 1):
-                return False
-            print("%s times ssh test......failed" % attempt)
-            time.sleep(2)
-    return False
+            # check whether the file is already existed
+            self.check_hostfile(HOST_FILE)
+        except Exception:
+            print("Generate default hostfile {0} under "
+                  "{1}".format(HOST_FILE, os.environ['HOME']))
+            self._generate_hostfile_via_installer()
+
+    def _generate_hostfile_via_installer(self):
+        self.hostfile = None
+
+        installer_type = str(os.environ['INSTALLER_TYPE'].lower())
+        installer_ip = str(os.environ['INSTALLER_IP'])
+
+        if installer_type not in ["fuel"]:
+            raise ValueError("{0} is not supported".format(installer_type))
+        if not installer_ip:
+            raise ValueError(
+                "The value of environment variable INSTALLER_IP is empty.")
+
+        cmd = "bash %s/generate_host_file.sh -t %s -i %s -d %s" % \
+              (SCRIPT_DIR, installer_type, installer_ip, HOST_FILE)
+        os.system(cmd)
+
+        self.hostfile = HOST_FILE
+
+    def fetch_host_ip_from_hostfile(self):
+        self.host_ip_list = []
+        print('Fetch host ips from hostfile...')
+        with open(self.hostfile, 'r') as f:
+            self.host_ip_list = re.findall('\d+.\d+.\d+.\d+', f.read())
+        if self.host_ip_list:
+            print("The remote compute nodes: {0}".format(self.host_ip_list))
+        else:
+            raise ValueError("The hostfile doesn't include host ip addresses.")
+
+    def check_hosts_ssh_connectivity(self):
+        results = map(lambda ip: self._ssh_is_ok(ip, self.keypair['private']),
+                      self.host_ip_list)
+        if not (len(results) == len(self.host_ip_list) and False not in results):
+            raise RuntimeError("Failed on checking hosts ssh connectivity.")
+
+    @staticmethod
+    def _ssh_is_ok(ip, private_key, attempts=100):
+        print('Check hosts {0} ssh connectivity...'.format(ip))
+        ssh = paramiko.SSHClient()
+        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+        ssh.connect(ip, key_filename=private_key)
+
+        for attempt in range(attempts):
+            try:
+                stdin, stdout, stderr = ssh.exec_command('uname')
+                if not stderr.readlines():
+                    print("{0}: SSH test successful.".format(ip))
+                    return True
+            except socket.error:
+                print("%s times ssh test......failed." % str(attempt + 1))
+                if attempt == (attempts - 1):
+                    return False
+                time.sleep(2)
+        return False
diff --git a/test-requirements.txt b/test-requirements.txt
index e434748e..cfbbdcdd 100644
--- a/test-requirements.txt
+++ b/test-requirements.txt
@@ -5,6 +5,7 @@
 tox
 pytest
 pytest-cov
+pytest-mock
 coverage
 pykwalify
 mock
diff --git a/tests/unit/util/env_test.py b/tests/unit/util/env_test.py
index 38ac988b..62d12a13 100644
--- a/tests/unit/util/env_test.py
+++ b/tests/unit/util/env_test.py
@@ -9,10 +9,40 @@
 
 import time
 
-import mock
 import pytest
+import mock
+from collections import defaultdict
+import socket
 
 from qtip.util import env
+from qtip.util.env import AnsibleEnvSetup
+
+
+@pytest.fixture(scope='session')
+def ansible_envsetup():
+    return AnsibleEnvSetup()
+
+
+@pytest.fixture()
+def hostfile(tmpdir):
+    fake_hostfile = tmpdir.join('hosts')
+    fake_hostfile.write("[hosts]\n")
+    fake_hostfile.write("10.20.0.3")
+    return fake_hostfile
+
+
+@pytest.fixture()
+def private_key(tmpdir):
+    fake_private_key = tmpdir.join('QtipKey')
+    fake_private_key.write("fake keypair")
+    return fake_private_key
+
+
+@pytest.fixture()
+def public_key(tmpdir):
+    fake_public_key = tmpdir.join('QtipKey.pub')
+    fake_public_key.write("fake public key")
+    return fake_public_key
 
 
 def test_all_files_exist(tmpdir):
@@ -35,60 +65,259 @@ def test_clean_file(tmpdir):
     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])
+def test_init(ansible_envsetup):
+    assert 'AnsibleEnvSetup' in str(type(ansible_envsetup))
+    assert ansible_envsetup.keypair == defaultdict(str)
+    assert ansible_envsetup.hostfile is None
+    assert ansible_envsetup.host_ip_list == []
 
-    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)
+def test_setup_exception(capsys, mocker, ansible_envsetup, hostfile):
+    with mock.patch.object(AnsibleEnvSetup, 'check_hostfile', side_effect=RuntimeError()):
+        mock_os = mocker.patch('sys.exit')
+        ansible_envsetup.setup({'hostfile': str(hostfile)})
+        out, error = capsys.readouterr()
+        assert out == '\n'
+        assert mock_os.call_count == 1
 
-    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)
+
+# TODO(zhihui_wu) Need find a smart way to write this pytest
+def test_setup(mocker, ansible_envsetup):
+    mock_check_hostfile = \
+        mocker.patch.object(AnsibleEnvSetup, 'check_hostfile')
+    mock_generate_default_hostfile = \
+        mocker.patch.object(AnsibleEnvSetup, 'generate_default_hostfile')
+    mock_fetch_ip = \
+        mocker.patch.object(AnsibleEnvSetup, 'fetch_host_ip_from_hostfile')
+    mock_check_keypair = \
+        mocker.patch.object(AnsibleEnvSetup, 'check_keypair')
+    mock_generate_default_keypair = \
+        mocker.patch.object(AnsibleEnvSetup, 'generate_default_keypair')
+    mock_pass_keypair = \
+        mocker.patch.object(AnsibleEnvSetup, 'pass_keypair_to_remote')
+    mock_check_ssh = \
+        mocker.patch.object(AnsibleEnvSetup, 'check_hosts_ssh_connectivity')
+
+    ansible_envsetup.setup({'keypair': str(private_key),
+                            'hostfile': str(hostfile)})
+    mock_check_hostfile.assert_called_with(str(hostfile))
+    mock_fetch_ip.assert_called_with()
+    mock_check_keypair.assert_called_with(str(private_key))
+    mock_pass_keypair.assert_called_with()
+    mock_check_ssh.assert_called_with()
+
+    ansible_envsetup.setup({'keypair': str(private_key)})
+    mock_generate_default_hostfile.assert_called_with()
+    mock_fetch_ip.assert_called_with()
+    mock_check_keypair.assert_called_with(str(private_key))
+    mock_pass_keypair.assert_called_with()
+    mock_check_ssh.assert_called_with()
+
+    ansible_envsetup.setup({'hostfile': str(hostfile)})
+    mock_check_hostfile.assert_called_with(str(hostfile))
+    mock_fetch_ip.assert_called_with()
+    mock_generate_default_keypair.assert_called_with()
+    mock_pass_keypair.assert_called_with()
+    mock_check_ssh.assert_called_with()
+
+    ansible_envsetup.setup()
+    mock_generate_default_hostfile.assert_called_with()
+    mock_fetch_ip.assert_called_with()
+    mock_generate_default_keypair.assert_called_with()
+    mock_pass_keypair.assert_called_with()
+    mock_check_ssh.assert_called_with()
+
+
+def test_check_keypair(mocker, ansible_envsetup, private_key, public_key):
+    with mocker.patch.object(env, 'all_files_exist', return_value=True):
+        ansible_envsetup.check_keypair(str(private_key))
+    assert ansible_envsetup.keypair['private'] == str(private_key)
+    assert ansible_envsetup.keypair['public'] == str(public_key)
+
+
+def test_check_keypair_failed(mocker, ansible_envsetup):
+    mocker.patch.object(env, 'all_files_exist', return_value=False)
+    with pytest.raises(RuntimeError) as excinfo:
+        ansible_envsetup.check_keypair(str(private_key))
+    assert 'The keypairs you in the configuration file ' \
+           'is invalid or not existed.' == str(excinfo.value)
+    assert ansible_envsetup.keypair['private'] == ''
+    assert ansible_envsetup.keypair['public'] == ''
+
+
+@pytest.mark.parametrize("file_existence, expected", [
+    (True, 0),
+    (False, 1)
+])
+def test_generate_default_keypair(mocker, ansible_envsetup, file_existence, expected):
+    mock_os = mocker.patch('os.system')
+    mocker.patch.object(env, 'all_files_exist', return_value=file_existence)
+    ansible_envsetup.generate_default_keypair()
+    assert mock_os.call_count == expected
+    assert ansible_envsetup.keypair['private'] == env.PRIVATE_KEY
+    assert ansible_envsetup.keypair['public'] == env.PUBLIC_KEY
+
+
+@pytest.mark.parametrize("ips, expected", [
+    (['10.20.0.3'], 1),
+    (['10.20.0.3', '10.20.0.4'], 2)
+])
+def test_pass_keypair_to_remote_successful(mocker, ansible_envsetup, ips, expected):
+    ansible_envsetup.host_ip_list = ips
+    mock_pass_keypair = \
+        mocker.patch.object(AnsibleEnvSetup, '_pass_keypair', return_value=True)
+    ansible_envsetup.pass_keypair_to_remote()
+    assert mock_pass_keypair.call_count == expected
+
+
+def test_pass_keypair_to_remote_failed(mocker, ansible_envsetup):
+    ansible_envsetup.host_ip_list = ['10.20.0.3']
+    mocker.patch.object(AnsibleEnvSetup, '_pass_keypair', return_value=False)
+    with pytest.raises(RuntimeError) as excinfo:
+        ansible_envsetup.pass_keypair_to_remote()
+    assert "Failed on passing keypair to remote." in str(excinfo.value)
 
 
-def test_generate_host_file(monkeypatch, tmpdir):
+def test_pass_keypair(monkeypatch, capsys, mocker, ansible_envsetup):
+    monkeypatch.setattr(time, 'sleep', lambda s: None)
+    mock_os = mocker.patch('os.system')
+    ansible_envsetup._pass_keypair('10.20.0.3', str(private_key))
+    assert mock_os.call_count == 2
+    out, error = capsys.readouterr()
+    assert "Pass keypair to remote hosts 10.20.0.3 successfully" in out
+
+
+def test_pass_keypair_exception(capsys, ansible_envsetup):
+    with mock.patch('os.system', side_effect=Exception()) as mock_os:
+        result = ansible_envsetup._pass_keypair('10.20.0.3', str(private_key))
+        assert result is False
+        out, error = capsys.readouterr()
+        assert out == '\n'
+        assert mock_os.call_count == 1
+
+
+def test_check_hostfile(mocker, ansible_envsetup, hostfile):
+    ansible_envsetup.check_hostfile(str(hostfile))
+    assert ansible_envsetup.hostfile == str(hostfile)
+
+    with pytest.raises(RuntimeError) as excinfo:
+        mocker.patch.object(env, 'all_files_exist', return_value=False)
+        ansible_envsetup.check_hostfile(str(hostfile))
+    assert str(excinfo.value) == 'The hostfile {0} is invalid or not ' \
+                                 'existed.'.format(str(hostfile))
+
+
+def test_default_hostfile_non_existed(mocker, ansible_envsetup):
+    with mocker.patch.object(env, 'all_files_exist', return_value=False):
+            mock_generate_hostfile_via_installer = \
+                mocker.patch.object(AnsibleEnvSetup,
+                                    '_generate_hostfile_via_installer')
+            ansible_envsetup.generate_default_hostfile()
+            mock_generate_hostfile_via_installer.assert_called_once_with()
+
+
+def test_default_hostfile_existed(mocker, ansible_envsetup):
+    with mocker.patch.object(env, 'all_files_exist', return_value=True):
+        mock_generate_hostfile_via_installer = \
+            mocker.patch.object(AnsibleEnvSetup,
+                                '_generate_hostfile_via_installer')
+        ansible_envsetup.generate_default_hostfile()
+        mock_generate_hostfile_via_installer.assert_not_called()
+
+
+@pytest.mark.parametrize("test_input, expected", [
+    (({}, KeyError), 'INSTALLER_TYPE'),
+    (({'INSTALLER_TYPE': 'fuel'}, KeyError), 'INSTALLER_IP'),
+    (({'INSTALLER_TYPE': 'fuel_1', 'INSTALLER_IP': '10.20.0.2'}, ValueError),
+     'fuel_1 is not supported'),
+    (({'INSTALLER_TYPE': 'fuel', 'INSTALLER_IP': ''}, ValueError),
+     'The value of environment variable INSTALLER_IP is empty')
+])
+def test_generate_hostfile_via_installer_exception(monkeypatch, ansible_envsetup, test_input, expected):
+    if test_input[0]:
+        for key in test_input[0]:
+            monkeypatch.setenv(key, test_input[0][key])
+
+    with pytest.raises(test_input[1]) as excinfo:
+        ansible_envsetup._generate_hostfile_via_installer()
+    assert expected in str(excinfo.value)
+
+
+def test_generate_hostfile_via_installer(monkeypatch, mocker, ansible_envsetup):
     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))
+    mock_os = mocker.patch('os.system')
+    ansible_envsetup._generate_hostfile_via_installer()
+    assert mock_os.call_count == 1
+    assert ansible_envsetup.hostfile == env.HOST_FILE
 
 
-def test_generate_keypair():
-    with mock.patch('os.system') as mock_os:
-        env.generate_keypair()
-        assert mock_os.call_count == 1
+def test_fetch_host_ip_from_hostfile(ansible_envsetup, hostfile):
+    ansible_envsetup.hostfile = str(hostfile)
+    ansible_envsetup.fetch_host_ip_from_hostfile()
+    assert ansible_envsetup.host_ip_list == ['10.20.0.3']
 
 
-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
+def test_fetch_host_ip_from_empty_hostfile(ansible_envsetup, tmpdir):
+    empty_hostfile = tmpdir.join('empty_hostfile')
+    empty_hostfile.write("")
+    ansible_envsetup.hostfile = str(empty_hostfile)
+    with pytest.raises(ValueError) as excinfo:
+        ansible_envsetup.fetch_host_ip_from_hostfile()
+    assert str(excinfo.value) == "The hostfile doesn't include host ip addresses."
+
+
+@pytest.mark.parametrize("ips, expected", [
+    (['10.20.0.3'], 1),
+    (['10.20.0.3', '10.20.0.4'], 2)
+])
+def test_check_hosts_ssh_connectivity(mocker, ansible_envsetup, ips, expected):
+    ansible_envsetup.host_ip_list = ips
+    mock_ssh_is_ok = \
+        mocker.patch.object(AnsibleEnvSetup, '_ssh_is_ok', return_value=True)
+    ansible_envsetup.check_hosts_ssh_connectivity()
+    assert mock_ssh_is_ok.call_count == expected
+
+
+def test_check_hosts_ssh_connectivity_failed(mocker, ansible_envsetup):
+    ansible_envsetup.host_ip_list = ['10.20.0.3']
+    mocker.patch.object(AnsibleEnvSetup, '_ssh_is_ok', return_value=False)
+    with pytest.raises(RuntimeError) as excinfo:
+        ansible_envsetup.check_hosts_ssh_connectivity()
+    assert "Failed on checking hosts ssh connectivity." == str(excinfo.value)
 
 
 @pytest.mark.parametrize("stderrinfo, expected", [
     ('', True),
     ('sorry', False)
 ])
-@mock.patch('paramiko.SSHClient')
-def test_ssh_is_ok(mock_sshclient, stderrinfo, expected):
+def test_ssh_is_ok(mocker, ansible_envsetup, private_key, stderrinfo, expected):
     stderr = mock.MagicMock()
     stderr.readlines.return_value = stderrinfo
+    mock_sshclient = mocker.patch('paramiko.SSHClient')
     test_ssh_client = mock_sshclient.return_value
     test_ssh_client.exec_command.return_value = ('', '', stderr)
-    result = env.ssh_is_ok('10.20.0.3')
+    result = ansible_envsetup._ssh_is_ok('10.20.0.3', str(private_key))
     assert result == expected
     test_ssh_client.connect.assert_called_once_with(
-        '10.20.0.3', key_filename=env.PRIVATE_KEY)
+        '10.20.0.3', key_filename=str(private_key))
     test_ssh_client.exec_command.assert_called_with('uname')
+
+
+@pytest.mark.parametrize("attempts, expected", [
+    (1,
+     'Check hosts 10.20.0.3 ssh connectivity...\n1 times ssh test......failed.\n'),
+    (2,
+     'Check hosts 10.20.0.3 ssh connectivity...\n'
+     '1 times ssh test......failed.\n'
+     '2 times ssh test......failed.\n')
+])
+def test_ssh_exception(capsys, monkeypatch, mocker, ansible_envsetup, attempts, expected):
+    monkeypatch.setattr(time, 'sleep', lambda s: None)
+    mock_sshclient = mocker.patch('paramiko.SSHClient')
+    test_ssh_client = mock_sshclient.return_value
+    test_ssh_client.exec_command.side_effect = socket.error()
+    result = ansible_envsetup._ssh_is_ok('10.20.0.3', str(private_key), attempts=attempts)
+    out, error = capsys.readouterr()
+    assert expected == out
+    assert result is False