add unit tests

Issue-ID: VNFSDK-350
Change-Id: Ib40c66476f5d8e32fbb4032596f67700551a963d
Signed-off-by: Moshe <moshehoa@amdocs.com>

4 files changed, 591 insertions, 0 deletions

diff --git a/vnftest/tests/unit/common/test_ssh.py b/vnftest/tests/unit/common/test_ssh.py
new file mode 100644
index 0000000..a571f29
--- /dev/null
+++ b/vnftest/tests/unit/common/test_ssh.py
@@ -0,0 +1,542 @@
+##############################################################################
+# Copyright 2018 EuropeanSoftwareMarketingLtd.
+# ===================================================================
+#  Licensed under the ApacheLicense, Version2.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
+#
+# 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
+##############################################################################
+# vnftest comment: this is a modified copy of
+# yardstick/tests/unit/test_ssh.py
+
+
+import os
+import socket
+import unittest
+from io import StringIO
+from itertools import count
+
+import mock
+from oslo_utils import encodeutils
+from vnftest import ssh
+
+from vnftest.ssh import SSHError, AutoConnectSSH, SSHTimeout, SSH
+
+
+class FakeParamikoException(Exception):
+    pass
+
+
+class SSHTestCase(unittest.TestCase):
+    """Test all small SSH methods."""
+
+    def setUp(self):
+        super(SSHTestCase, self).setUp()
+        self.test_client = ssh.SSH("root", "example.net")
+
+    @mock.patch("vnftest.ssh.SSH._get_pkey")
+    def test_construct(self, mock_ssh__get_pkey):
+        mock_ssh__get_pkey.return_value = "pkey"
+        test_ssh = ssh.SSH("root", "example.net", port=33, pkey="key",
+                           key_filename="kf", password="secret")
+        mock_ssh__get_pkey.assert_called_once_with("key")
+        self.assertEqual("root", test_ssh.user)
+        self.assertEqual("example.net", test_ssh.host)
+        self.assertEqual(33, test_ssh.port)
+        self.assertEqual("pkey", test_ssh.pkey)
+        self.assertEqual("kf", test_ssh.key_filename)
+        self.assertEqual("secret", test_ssh.password)
+
+    @mock.patch("vnftest.ssh.SSH._get_pkey")
+    def test_ssh_from_node(self, mock_ssh__get_pkey):
+        mock_ssh__get_pkey.return_value = "pkey"
+        node = {
+            "user": "root", "ip": "example.net", "ssh_port": 33,
+            "key_filename": "kf", "password": "secret"
+        }
+        test_ssh = ssh.SSH.from_node(node)
+        self.assertEqual("root", test_ssh.user)
+        self.assertEqual("example.net", test_ssh.host)
+        self.assertEqual(33, test_ssh.port)
+        self.assertEqual("kf", test_ssh.key_filename)
+        self.assertEqual("secret", test_ssh.password)
+
+    @mock.patch("vnftest.ssh.SSH._get_pkey")
+    def test_ssh_from_node_password_default(self, mock_ssh__get_pkey):
+        mock_ssh__get_pkey.return_value = "pkey"
+        node = {
+            "user": "root", "ip": "example.net", "ssh_port": 33,
+            "key_filename": "kf"
+        }
+        test_ssh = ssh.SSH.from_node(node)
+        self.assertEqual("root", test_ssh.user)
+        self.assertEqual("example.net", test_ssh.host)
+        self.assertEqual(33, test_ssh.port)
+        self.assertEqual("kf", test_ssh.key_filename)
+        self.assertIsNone(test_ssh.password)
+
+    @mock.patch("vnftest.ssh.SSH._get_pkey")
+    def test_ssh_from_node_ssh_port_default(self, mock_ssh__get_pkey):
+        mock_ssh__get_pkey.return_value = "pkey"
+        node = {
+            "user": "root", "ip": "example.net",
+            "key_filename": "kf", "password": "secret"
+        }
+        test_ssh = ssh.SSH.from_node(node)
+        self.assertEqual("root", test_ssh.user)
+        self.assertEqual("example.net", test_ssh.host)
+        self.assertEqual(ssh.SSH.SSH_PORT, test_ssh.port)
+        self.assertEqual("kf", test_ssh.key_filename)
+        self.assertEqual("secret", test_ssh.password)
+
+    @mock.patch("vnftest.ssh.SSH._get_pkey")
+    def test_ssh_from_node_key_filename_default(self, mock_ssh__get_pkey):
+        mock_ssh__get_pkey.return_value = "pkey"
+        node = {
+            "user": "root", "ip": "example.net", "ssh_port": 33,
+            "password": "secret"
+        }
+        test_ssh = ssh.SSH.from_node(node)
+        self.assertEqual("root", test_ssh.user)
+        self.assertEqual("example.net", test_ssh.host)
+        self.assertEqual(33, test_ssh.port)
+        self.assertIsNone(test_ssh.key_filename)
+        self.assertEqual("secret", test_ssh.password)
+
+    def test_construct_default(self):
+        self.assertEqual("root", self.test_client.user)
+        self.assertEqual("example.net", self.test_client.host)
+        self.assertEqual(22, self.test_client.port)
+        self.assertIsNone(self.test_client.pkey)
+        self.assertIsNone(self.test_client.key_filename)
+        self.assertIsNone(self.test_client.password)
+
+    @mock.patch("vnftest.ssh.paramiko")
+    def test__get_pkey_invalid(self, mock_paramiko):
+        mock_paramiko.SSHException = FakeParamikoException
+        rsa = mock_paramiko.rsakey.RSAKey
+        dss = mock_paramiko.dsskey.DSSKey
+        rsa.from_private_key.side_effect = mock_paramiko.SSHException
+        dss.from_private_key.side_effect = mock_paramiko.SSHException
+        self.assertRaises(SSHError, self.test_client._get_pkey, "key")
+
+    @mock.patch("vnftest.ssh.six.moves.StringIO")
+    @mock.patch("vnftest.ssh.paramiko")
+    def test__get_pkey_dss(self, mock_paramiko, mock_string_io):
+        mock_paramiko.SSHException = FakeParamikoException
+        mock_string_io.return_value = "string_key"
+        mock_paramiko.dsskey.DSSKey.from_private_key.return_value = "dss_key"
+        rsa = mock_paramiko.rsakey.RSAKey
+        rsa.from_private_key.side_effect = mock_paramiko.SSHException
+        key = self.test_client._get_pkey("key")
+        dss_calls = mock_paramiko.dsskey.DSSKey.from_private_key.mock_calls
+        self.assertEqual([mock.call("string_key")], dss_calls)
+        self.assertEqual(key, "dss_key")
+        mock_string_io.assert_called_once_with("key")
+
+    @mock.patch("vnftest.ssh.six.moves.StringIO")
+    @mock.patch("vnftest.ssh.paramiko")
+    def test__get_pkey_rsa(self, mock_paramiko, mock_string_io):
+        mock_paramiko.SSHException = FakeParamikoException
+        mock_string_io.return_value = "string_key"
+        mock_paramiko.rsakey.RSAKey.from_private_key.return_value = "rsa_key"
+        dss = mock_paramiko.dsskey.DSSKey
+        dss.from_private_key.side_effect = mock_paramiko.SSHException
+        key = self.test_client._get_pkey("key")
+        rsa_calls = mock_paramiko.rsakey.RSAKey.from_private_key.mock_calls
+        self.assertEqual([mock.call("string_key")], rsa_calls)
+        self.assertEqual(key, "rsa_key")
+        mock_string_io.assert_called_once_with("key")
+
+    @mock.patch("vnftest.ssh.SSH._get_pkey")
+    @mock.patch("vnftest.ssh.paramiko")
+    def test__get_client(self, mock_paramiko, mock_ssh__get_pkey):
+        mock_ssh__get_pkey.return_value = "key"
+        fake_client = mock.Mock()
+        mock_paramiko.SSHClient.return_value = fake_client
+        mock_paramiko.AutoAddPolicy.return_value = "autoadd"
+
+        test_ssh = ssh.SSH("admin", "example.net", pkey="key")
+        client = test_ssh._get_client()
+
+        self.assertEqual(fake_client, client)
+        client_calls = [
+            mock.call.set_missing_host_key_policy("autoadd"),
+            mock.call.connect("example.net", username="admin",
+                              port=22, pkey="key", key_filename=None,
+                              password=None,
+                              allow_agent=False, look_for_keys=False,
+                              timeout=1),
+        ]
+        self.assertEqual(client_calls, client.mock_calls)
+
+    @mock.patch("vnftest.ssh.SSH._get_pkey")
+    @mock.patch("vnftest.ssh.paramiko")
+    def test__get_client_with_exception(self, mock_paramiko, mock_ssh__get_pkey):
+        class MyError(Exception):
+            pass
+
+        mock_ssh__get_pkey.return_value = "pkey"
+        fake_client = mock.Mock()
+        fake_client.connect.side_effect = MyError
+        fake_client.set_missing_host_key_policy.return_value = None
+        mock_paramiko.SSHClient.return_value = fake_client
+        mock_paramiko.AutoAddPolicy.return_value = "autoadd"
+
+        test_ssh = ssh.SSH("admin", "example.net", pkey="key")
+
+        with self.assertRaises(SSHError) as raised:
+            test_ssh._get_client()
+
+        mock_paramiko.SSHClient.assert_called_once()
+        mock_paramiko.AutoAddPolicy.assert_called_once()
+        fake_client.set_missing_host_key_policy.assert_called_once()
+        fake_client.connect.assert_called_once()
+        exc_str = str(raised.exception)
+        self.assertIn('raised during connect', exc_str)
+        self.assertIn('MyError', exc_str)
+
+    @mock.patch("vnftest.ssh.SSH._get_pkey")
+    @mock.patch("vnftest.ssh.paramiko")
+    def test_copy(self, mock_paramiko, mock_ssh__get_pkey):
+        mock_ssh__get_pkey.return_value = "pkey"
+        fake_client = mock.Mock()
+        fake_client.connect.side_effect = IOError
+        mock_paramiko.SSHClient.return_value = fake_client
+        mock_paramiko.AutoAddPolicy.return_value = "autoadd"
+
+        test_ssh = ssh.SSH("admin", "example.net", pkey="key")
+        result = test_ssh.copy()
+        self.assertIsNot(test_ssh, result)
+
+    def test_close(self):
+        with mock.patch.object(self.test_client, "_client") as m_client:
+            self.test_client.close()
+        m_client.close.assert_called_once_with()
+        self.assertFalse(self.test_client._client)
+
+    @mock.patch("vnftest.ssh.time")
+    def test_wait_timeout(self, mock_time):
+        mock_time.time.side_effect = [1, 50, 150]
+        self.test_client.execute = mock.Mock(side_effect=[SSHError,
+                                                          SSHError,
+                                                          0])
+        self.assertRaises(SSHTimeout, self.test_client.wait)
+        self.assertEqual([mock.call("uname")] * 2,
+                         self.test_client.execute.mock_calls)
+
+    @mock.patch("vnftest.ssh.time")
+    def test_wait(self, mock_time):
+        mock_time.time.side_effect = [1, 50, 100]
+        self.test_client.execute = mock.Mock(side_effect=[SSHError,
+                                                          SSHError,
+                                                          0])
+        self.test_client.wait()
+        self.assertEqual([mock.call("uname")] * 3,
+                         self.test_client.execute.mock_calls)
+
+    @mock.patch("vnftest.ssh.paramiko")
+    def test_send_command(self, _):
+        paramiko_sshclient = self.test_client._get_client()
+        with mock.patch.object(paramiko_sshclient, "exec_command") \
+                as mock_paramiko_exec_command:
+            self.test_client.send_command('cmd')
+        mock_paramiko_exec_command.assert_called_once_with('cmd',
+                                                           get_pty=True)
+
+
+class SSHRunTestCase(unittest.TestCase):
+    """Test SSH.run method in different aspects.
+
+    Also tested method "execute".
+    """
+
+    def setUp(self):
+        super(SSHRunTestCase, self).setUp()
+
+        self.fake_client = mock.Mock()
+        self.fake_session = mock.Mock()
+        self.fake_transport = mock.Mock()
+
+        self.fake_transport.open_session.return_value = self.fake_session
+        self.fake_client.get_transport.return_value = self.fake_transport
+
+        self.fake_session.recv_ready.return_value = False
+        self.fake_session.recv_stderr_ready.return_value = False
+        self.fake_session.send_ready.return_value = False
+        self.fake_session.exit_status_ready.return_value = True
+        self.fake_session.recv_exit_status.return_value = 0
+
+        self.test_client = ssh.SSH("admin", "example.net")
+        self.test_client._get_client = mock.Mock(return_value=self.fake_client)
+
+    @mock.patch("vnftest.ssh.select")
+    def test_execute(self, mock_select):
+        mock_select.select.return_value = ([], [], [])
+        self.fake_session.recv_ready.side_effect = [1, 0, 0]
+        self.fake_session.recv_stderr_ready.side_effect = [1, 0]
+        self.fake_session.recv.return_value = "ok"
+        self.fake_session.recv_stderr.return_value = "error"
+        self.fake_session.exit_status_ready.return_value = 1
+        self.fake_session.recv_exit_status.return_value = 127
+        self.assertEqual((127, "ok", "error"), self.test_client.execute("cmd"))
+        self.fake_session.exec_command.assert_called_once_with("cmd")
+
+    @mock.patch("vnftest.ssh.select")
+    def test_execute_args(self, mock_select):
+        mock_select.select.return_value = ([], [], [])
+        self.fake_session.recv_ready.side_effect = [1, 0, 0]
+        self.fake_session.recv_stderr_ready.side_effect = [1, 0]
+        self.fake_session.recv.return_value = "ok"
+        self.fake_session.recv_stderr.return_value = "error"
+        self.fake_session.exit_status_ready.return_value = 1
+        self.fake_session.recv_exit_status.return_value = 127
+
+        result = self.test_client.execute("cmd arg1 'arg2 with space'")
+        self.assertEqual((127, "ok", "error"), result)
+        self.fake_session.exec_command.assert_called_once_with(
+            "cmd arg1 'arg2 with space'")
+
+    @mock.patch("vnftest.ssh.select")
+    def test_run(self, mock_select):
+        mock_select.select.return_value = ([], [], [])
+        self.assertEqual(0, self.test_client.run("cmd"))
+
+    @mock.patch("vnftest.ssh.select")
+    def test_run_nonzero_status(self, mock_select):
+        mock_select.select.return_value = ([], [], [])
+        self.fake_session.recv_exit_status.return_value = 1
+        self.assertRaises(SSHError, self.test_client.run, "cmd")
+        self.assertEqual(1, self.test_client.run("cmd", raise_on_error=False))
+
+    @mock.patch("vnftest.ssh.select")
+    def test_run_stdout(self, mock_select):
+        mock_select.select.return_value = ([], [], [])
+        self.fake_session.recv_ready.side_effect = [True, True, False]
+        self.fake_session.recv.side_effect = ["ok1", "ok2"]
+        stdout = mock.Mock()
+        self.test_client.run("cmd", stdout=stdout)
+        self.assertEqual([mock.call("ok1"), mock.call("ok2")],
+                         stdout.write.mock_calls)
+
+    @mock.patch("vnftest.ssh.select")
+    def test_run_stderr(self, mock_select):
+        mock_select.select.return_value = ([], [], [])
+        self.fake_session.recv_stderr_ready.side_effect = [True, False]
+        self.fake_session.recv_stderr.return_value = "error"
+        stderr = mock.Mock()
+        self.test_client.run("cmd", stderr=stderr)
+        stderr.write.assert_called_once_with("error")
+
+    @mock.patch("vnftest.ssh.select")
+    def test_run_stdin(self, mock_select):
+        """Test run method with stdin.
+
+        Third send call was called with "e2" because only 3 bytes was sent
+        by second call. So remainig 2 bytes of "line2" was sent by third call.
+        """
+        mock_select.select.return_value = ([], [], [])
+        self.fake_session.exit_status_ready.side_effect = [0, 0, 0, True]
+        self.fake_session.send_ready.return_value = True
+        self.fake_session.send.side_effect = [5, 3, 2]
+        fake_stdin = mock.Mock()
+        fake_stdin.read.side_effect = ["line1", "line2", ""]
+        fake_stdin.closed = False
+
+        def close():
+            fake_stdin.closed = True
+        fake_stdin.close = mock.Mock(side_effect=close)
+        self.test_client.run("cmd", stdin=fake_stdin)
+        call = mock.call
+        send_calls = [call(encodeutils.safe_encode("line1", "utf-8")),
+                      call(encodeutils.safe_encode("line2", "utf-8")),
+                      call(encodeutils.safe_encode("e2", "utf-8"))]
+        self.assertEqual(send_calls, self.fake_session.send.mock_calls)
+
+    @mock.patch("vnftest.ssh.select")
+    def test_run_stdin_keep_open(self, mock_select):
+        """Test run method with stdin.
+
+        Third send call was called with "e2" because only 3 bytes was sent
+        by second call. So remainig 2 bytes of "line2" was sent by third call.
+        """
+        mock_select.select.return_value = ([], [], [])
+        self.fake_session.exit_status_ready.side_effect = [0, 0, 0, True]
+        self.fake_session.send_ready.return_value = True
+        self.fake_session.send.side_effect = len
+        fake_stdin = StringIO(u"line1\nline2\n")
+        self.test_client.run("cmd", stdin=fake_stdin, keep_stdin_open=True)
+        call = mock.call
+        send_calls = [call(encodeutils.safe_encode("line1\nline2\n", "utf-8"))]
+        self.assertEqual(send_calls, self.fake_session.send.mock_calls)
+
+    @mock.patch("vnftest.ssh.select")
+    def test_run_select_error(self, mock_select):
+        self.fake_session.exit_status_ready.return_value = False
+        mock_select.select.return_value = ([], [], [True])
+        self.assertRaises(SSHError, self.test_client.run, "cmd")
+
+    @mock.patch("vnftest.ssh.time")
+    @mock.patch("vnftest.ssh.select")
+    def test_run_timemout(self, mock_select, mock_time):
+        mock_time.time.side_effect = [1, 3700]
+        mock_select.select.return_value = ([], [], [])
+        self.fake_session.exit_status_ready.return_value = False
+        self.assertRaises(SSHTimeout, self.test_client.run, "cmd")
+
+    @mock.patch("vnftest.ssh.open", create=True)
+    def test__put_file_shell(self, mock_open):
+        with mock.patch.object(self.test_client, "run") as run_mock:
+            self.test_client._put_file_shell("localfile", "remotefile", 0o42)
+            run_mock.assert_called_once_with(
+                'cat > "remotefile"&& chmod -- 042 "remotefile"',
+                stdin=mock_open.return_value.__enter__.return_value)
+
+    @mock.patch("vnftest.ssh.open", create=True)
+    def test__put_file_shell_space(self, mock_open):
+        with mock.patch.object(self.test_client, "run") as run_mock:
+            self.test_client._put_file_shell("localfile",
+                                             "filename with space", 0o42)
+            run_mock.assert_called_once_with(
+                'cat > "filename with space"&& chmod -- 042 "filename with '
+                'space"',
+                stdin=mock_open.return_value.__enter__.return_value)
+
+    @mock.patch("vnftest.ssh.open", create=True)
+    def test__put_file_shell_tilde(self, mock_open):
+        with mock.patch.object(self.test_client, "run") as run_mock:
+            self.test_client._put_file_shell("localfile", "~/remotefile", 0o42)
+            run_mock.assert_called_once_with(
+                'cat > ~/"remotefile"&& chmod -- 042 ~/"remotefile"',
+                stdin=mock_open.return_value.__enter__.return_value)
+
+    @mock.patch("vnftest.ssh.open", create=True)
+    def test__put_file_shell_tilde_spaces(self, mock_open):
+        with mock.patch.object(self.test_client, "run") as run_mock:
+            self.test_client._put_file_shell("localfile", "~/file with space",
+                                             0o42)
+            run_mock.assert_called_once_with(
+                'cat > ~/"file with space"&& chmod -- 042 ~/"file with space"',
+                stdin=mock_open.return_value.__enter__.return_value)
+
+    @mock.patch("vnftest.ssh.os.stat")
+    def test__put_file_sftp(self, mock_stat):
+        sftp = self.fake_client.open_sftp.return_value = mock.MagicMock()
+        sftp.__enter__.return_value = sftp
+
+        mock_stat.return_value = os.stat_result([0o753] + [0] * 9)
+
+        self.test_client._put_file_sftp("localfile", "remotefile")
+
+        sftp.put.assert_called_once_with("localfile", "remotefile")
+        mock_stat.assert_any_call("localfile")
+        sftp.chmod.assert_any_call("remotefile", 0o753)
+        sftp.__exit__.assert_called_once_with(None, None, None)
+
+    def test__put_file_sftp_mode(self):
+        sftp = self.fake_client.open_sftp.return_value = mock.MagicMock()
+        sftp.__enter__.return_value = sftp
+
+        self.test_client._put_file_sftp("localfile", "remotefile", mode=0o753)
+
+        sftp.put.assert_called_once_with("localfile", "remotefile")
+        sftp.chmod.assert_called_once_with("remotefile", 0o753)
+        sftp.__exit__.assert_called_once_with(None, None, None)
+
+    def test_put_file_SSHException(self):
+        exc = ssh.paramiko.SSHException
+        self.test_client._put_file_sftp = mock.Mock(side_effect=exc())
+        self.test_client._put_file_shell = mock.Mock()
+
+        self.test_client.put_file("foo", "bar", 42)
+        self.test_client._put_file_sftp.assert_called_once_with("foo", "bar",
+                                                                mode=42)
+        self.test_client._put_file_shell.assert_called_once_with("foo", "bar",
+                                                                 mode=42)
+
+    def test_put_file_socket_error(self):
+        exc = socket.error
+        self.test_client._put_file_sftp = mock.Mock(side_effect=exc())
+        self.test_client._put_file_shell = mock.Mock()
+
+        self.test_client.put_file("foo", "bar", 42)
+        self.test_client._put_file_sftp.assert_called_once_with("foo", "bar",
+                                                                mode=42)
+        self.test_client._put_file_shell.assert_called_once_with("foo", "bar",
+                                                                 mode=42)
+
+    @mock.patch("vnftest.ssh.os.stat")
+    def test_put_file_obj_with_mode(self, mock_stat):
+        sftp = self.fake_client.open_sftp.return_value = mock.MagicMock()
+        sftp.__enter__.return_value = sftp
+
+        mock_stat.return_value = os.stat_result([0o753] + [0] * 9)
+
+        self.test_client.put_file_obj("localfile", "remotefile", 'my_mode')
+
+        sftp.__enter__.assert_called_once()
+        sftp.putfo.assert_called_once_with("localfile", "remotefile")
+        sftp.chmod.assert_called_once_with("remotefile", 'my_mode')
+        sftp.__exit__.assert_called_once_with(None, None, None)
+
+
+class TestAutoConnectSSH(unittest.TestCase):
+
+    def test__connect_loop(self):
+        auto_connect_ssh = AutoConnectSSH('user1', 'host1', wait=0)
+        auto_connect_ssh._get_client = mock__get_client = mock.Mock()
+
+        auto_connect_ssh._connect()
+        mock__get_client.assert_called_once()
+
+    def test_get_class(self):
+        auto_connect_ssh = AutoConnectSSH('user1', 'host1')
+
+        self.assertEqual(auto_connect_ssh.get_class(), AutoConnectSSH)
+
+    def test_drop_connection(self):
+        auto_connect_ssh = AutoConnectSSH('user1', 'host1')
+        self.assertFalse(auto_connect_ssh._client)
+        auto_connect_ssh._client = True
+        auto_connect_ssh.drop_connection()
+        self.assertFalse(auto_connect_ssh._client)
+
+    @mock.patch('vnftest.ssh.SCPClient')
+    def test_put(self, mock_scp_client_type):
+        auto_connect_ssh = AutoConnectSSH('user1', 'host1')
+        auto_connect_ssh._client = mock.Mock()
+
+        auto_connect_ssh.put('a', 'z')
+        with mock_scp_client_type() as mock_scp_client:
+            mock_scp_client.put.assert_called_once()
+
+    @mock.patch('vnftest.ssh.SCPClient')
+    def test_get(self, mock_scp_client_type):
+        auto_connect_ssh = AutoConnectSSH('user1', 'host1')
+        auto_connect_ssh._client = mock.Mock()
+
+        auto_connect_ssh.get('a', 'z')
+        with mock_scp_client_type() as mock_scp_client:
+            mock_scp_client.get.assert_called_once()
+
+    def test_put_file(self):
+        auto_connect_ssh = AutoConnectSSH('user1', 'host1')
+        auto_connect_ssh._client = mock.Mock()
+        auto_connect_ssh._put_file_sftp = mock_put_sftp = mock.Mock()
+
+        auto_connect_ssh.put_file('a', 'b')
+        mock_put_sftp.assert_called_once()
+
+    def test_execute(self):
+        auto_connect_ssh = AutoConnectSSH('user1', 'host1')
+        auto_connect_ssh._client = mock.Mock()
+        auto_connect_ssh.run = mock.Mock(return_value=0)
+        exit_code, _, _ = auto_connect_ssh.execute('')
+        self.assertEqual(exit_code, 0)
+
+
diff --git a/vnftest/tests/unit/context/__init__.py b/vnftest/tests/unit/context/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/vnftest/tests/unit/context/__init__.py
diff --git a/vnftest/tests/unit/context/test_heat.py b/vnftest/tests/unit/context/test_heat.py
new file mode 100644
index 0000000..89799b6
--- /dev/null
+++ b/vnftest/tests/unit/context/test_heat.py
@@ -0,0 +1,47 @@
+##############################################################################
+# Copyright 2018 EuropeanSoftwareMarketingLtd.
+# ===================================================================
+#  Licensed under the ApacheLicense, Version2.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
+#
+# 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 unittest
+
+import mock
+import os
+
+from vnftest.contexts.heat import HeatContext
+
+from vnftest.common import constants as consts
+from vnftest.core import task
+from vnftest.common import openstack_utils
+
+
+class HeatTestCase(unittest.TestCase):
+
+    @mock.patch.object(HeatContext, 'check_environment')
+    @mock.patch.object(HeatContext, '_create_new_stack')
+    @mock.patch.object(HeatContext, 'get_neutron_info')
+    @mock.patch.object(openstack_utils, 'get_shade_client')
+    @mock.patch.object(openstack_utils, 'get_shade_operator_client')
+    def test_heat(self, mock_check_env, mock_create, mocke_neutron, shade_client, operator_client):
+        params = {
+            "task_id": "123",
+            "name": "heat-test",
+            "image": "test_image",
+            "flavor" : "test_flavor",
+            "user": "test_user",
+            "servers": {}}
+        mock_create.return_value = {}
+        h = HeatContext()
+        h.init(params)
+        h.deploy()
+        h._get_server("dummy")
+        h.undeploy()
diff --git a/vnftest/tests/unit/core/test_task.py b/vnftest/tests/unit/core/test_task.py
index 7f26d61..abec5ae 100644
--- a/vnftest/tests/unit/core/test_task.py
+++ b/vnftest/tests/unit/core/test_task.py
@@ -67,6 +67,8 @@ class TaskTestCase(unittest.TestCase):
         mock_base_runner.Runner.get.return_value = runner
         t._run([step], 'dummy_case', False, "vnftest.out", {})
         self.assertTrue(runner.run.called)
+        results = t.task_info.result()
+        self.assertEqual(results["testcases"][0]["status"], "FINISHED")
 
     def test_parse_suite_no_constraint_no_args(self):
         SAMPLE_step_PATH = "no_constraint_no_args_step_sample.yaml"