From 27cc6c9a5ce4092ff03315aaaff6cb2025e16204 Mon Sep 17 00:00:00 2001
From: Jon Bernard <jobernar@redhat.com>
Date: Mon, 23 Nov 2015 17:41:43 -0500
Subject: [PATCH] Allow RBDClient to be used from a with-statement

This allows users of the RBD client to use a 'with' statement and bound
the scope of a client connection to a python block, which makes for
lovely looking code under the right conditions.

Change-Id: I455e8b6063dff36794d8aec7922ee1f26b3cf298
---
 os_brick/initiator/linuxrbd.py            |  6 +++++
 os_brick/tests/initiator/test_linuxrbd.py | 32 +++++++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/os_brick/initiator/linuxrbd.py b/os_brick/initiator/linuxrbd.py
index da2019be3..cd939bc0f 100644
--- a/os_brick/initiator/linuxrbd.py
+++ b/os_brick/initiator/linuxrbd.py
@@ -57,6 +57,12 @@ class RBDClient(object):
 
         self.client, self.ioctx = self.connect()
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type_, value, traceback):
+        self.disconnect()
+
     def connect(self):
         client = self.rados.Rados(rados_id=self.rbd_user,
                                   conffile=self.rbd_conf)
diff --git a/os_brick/tests/initiator/test_linuxrbd.py b/os_brick/tests/initiator/test_linuxrbd.py
index 2a274f53f..ed33f56de 100644
--- a/os_brick/tests/initiator/test_linuxrbd.py
+++ b/os_brick/tests/initiator/test_linuxrbd.py
@@ -14,6 +14,38 @@ import mock
 
 from os_brick.initiator import linuxrbd
 from os_brick.tests import base
+from oslo_utils import encodeutils
+
+
+class RBDClientTestCase(base.TestCase):
+
+    def setUp(self):
+        super(RBDClientTestCase, self).setUp()
+
+    @mock.patch('os_brick.initiator.linuxrbd.rbd')
+    @mock.patch('os_brick.initiator.linuxrbd.rados')
+    def test_with_client(self, mock_rados, mock_rbd):
+        with linuxrbd.RBDClient('test_user', 'test_pool') as client:
+
+            # Verify object attributes are assigned as expected
+            self.assertEqual('/etc/ceph/ceph.conf', client.rbd_conf)
+            self.assertEqual(encodeutils.safe_encode('test_user'),
+                             client.rbd_user)
+            self.assertEqual(encodeutils.safe_encode('test_pool'),
+                             client.rbd_pool)
+
+            # Assert connect is called with correct paramaters
+            mock_rados.Rados.assert_called_once_with(
+                rados_id=encodeutils.safe_encode('test_user'),
+                conffile='/etc/ceph/ceph.conf')
+
+            # Ensure correct calls to connect to cluster
+            self.assertEqual(
+                1, mock_rados.Rados.return_value.connect.call_count)
+            mock_rados.Rados.return_value.open_ioctx.assert_called_once_with(
+                encodeutils.safe_encode('test_pool'))
+
+        self.assertEqual(1, mock_rados.Rados.return_value.shutdown.call_count)
 
 
 class RBDVolumeIOWrapperTestCase(base.TestCase):