From bd9ac1c521be43943395bcd108b1bd12bb7c2403 Mon Sep 17 00:00:00 2001
From: Hadi Esiely <hadi.esiely-barrera@jhuapl.edu>
Date: Tue, 9 Feb 2016 14:13:22 -0500
Subject: [PATCH] Server Failover Unit Tests

This pull request adds some unit tests for the server failover code that was recently merged.

Signed-off-by: Hadi Esiely <hadi.esiely-barrera@jhuapl.edu>
---
 kmip/pykmip.conf                             |  2 +-
 kmip/services/kmip_client.py                 |  3 +-
 kmip/tests/unit/services/test_kmip_client.py | 62 ++++++++++++++------
 3 files changed, 48 insertions(+), 19 deletions(-)

diff --git a/kmip/pykmip.conf b/kmip/pykmip.conf
index 2dae375..1ef9024 100644
--- a/kmip/pykmip.conf
+++ b/kmip/pykmip.conf
@@ -5,7 +5,7 @@ keyfile=None
 certfile=None
 cert_reqs=CERT_REQUIRED
 ssl_version=PROTOCOL_SSLv23
-ca_certs=./demos/certs/server.crt
+ca_certs=./kmip/demos/certs/server.crt
 do_handshake_on_connect=True
 suppress_ragged_eofs=True
 username=None
diff --git a/kmip/services/kmip_client.py b/kmip/services/kmip_client.py
index 94f9e1e..67362a7 100644
--- a/kmip/services/kmip_client.py
+++ b/kmip/services/kmip_client.py
@@ -220,9 +220,10 @@ class KMIPProxy(KMIP):
                 self.logger.error("An error occurred while connecting to "
                                   "appliance " + self.host)
                 self.socket.close()
-                self.socket = None
             else:
                 return
+
+        self.socket = None
         raise e
 
     def _create_socket(self, sock):
diff --git a/kmip/tests/unit/services/test_kmip_client.py b/kmip/tests/unit/services/test_kmip_client.py
index 3f4cf48..c400f30 100644
--- a/kmip/tests/unit/services/test_kmip_client.py
+++ b/kmip/tests/unit/services/test_kmip_client.py
@@ -62,8 +62,9 @@ from kmip.services.results import RekeyKeyPairResult
 import kmip.core.utils as utils
 
 import mock
-
+import os
 import socket
+import ssl
 
 
 class TestKMIPClient(TestCase):
@@ -77,6 +78,16 @@ class TestKMIPClient(TestCase):
 
         self.client = KMIPProxy()
 
+        KMIP_PORT = 9090
+        CA_CERTS_PATH = os.path.normpath(os.path.join(os.path.dirname(
+            os.path.abspath(__file__)), '../utils/certs/server.crt'))
+
+        self.mock_client = KMIPProxy(host="IP_ADDR_1, IP_ADDR_2",
+                                     port=KMIP_PORT, ca_certs=CA_CERTS_PATH)
+        self.mock_client.socket = mock.MagicMock()
+        self.mock_client.socket.connect = mock.MagicMock()
+        self.mock_client.socket.close = mock.MagicMock()
+
     def tearDown(self):
         super(TestKMIPClient, self).tearDown()
 
@@ -541,25 +552,42 @@ class TestKMIPClient(TestCase):
         self.assertRaises(expected_error, self.client._set_variables,
                           **kwargs)
 
-    @mock.patch('socket.socket.connect')
-    @mock.patch('ssl.SSLSocket.gettimeout')
-    def test_timeout_all_hosts(self, mock_ssl_timeout, mock_connect_return):
+    @mock.patch.object(KMIPProxy, '_create_socket')
+    def test_open_server_conn_failover_fail(self, mock_create_socket):
         """
-        This test verifies that the client will throw an exception if no
-        hosts are available for connection.
+        This test verifies that the KMIP client throws an exception if no
+        servers are available for connection
         """
+        mock_create_socket.return_value = mock.MagicMock()
 
-        mock_ssl_timeout.return_value = 1
-        mock_connect_return.return_value = socket.timeout
-        try:
-            self.client.open()
-        except Exception as e:
-            # TODO: once the exception is properly defined in the
-            # kmip_client.py file this test needs to change to reflect that.
-            self.assertIsInstance(e, Exception)
-            self.client.close()
-        else:
-            self.client.close()
+        # Assumes both IP addresses fail connection attempts
+        self.mock_client.socket.connect.side_effect = [Exception, Exception]
+
+        self.assertRaises(Exception, self.mock_client.open)
+
+    @mock.patch.object(KMIPProxy, '_create_socket')
+    def test_open_server_conn_failover_succeed(self, mock_create_socket):
+        """
+        This test verifies that the KMIP client can setup a connection if at
+        least one connection is established
+        """
+        mock_create_socket.return_value = mock.MagicMock()
+
+        # Assumes IP_ADDR_1 is a bad address and IP_ADDR_2 is a good address
+        self.mock_client.socket.connect.side_effect = [Exception, None]
+
+        self.mock_client.open()
+
+        self.assertEqual('IP_ADDR_2', self.mock_client.host)
+
+    def test_socket_ssl_wrap(self):
+        """
+        This test tests that the KMIP socket is successfully wrapped into an
+        ssl socket
+        """
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        self.client._create_socket(sock)
+        self.assertEqual(ssl.SSLSocket, type(self.client.socket))
 
 
 class TestClientProfileInformation(TestCase):