3 files changed, 21 insertions, 30 deletions

diff --git a/src/test/librbd/managed_lock/test_mock_BreakRequest.cc b/src/test/librbd/managed_lock/test_mock_BreakRequest.cc
index 3ebc233024f..c1dff856a0d 100644
--- a/src/test/librbd/managed_lock/test_mock_BreakRequest.cc
+++ b/src/test/librbd/managed_lock/test_mock_BreakRequest.cc
@@ -62,7 +62,8 @@ public:
   }
 
   void expect_blacklist_add(MockTestImageCtx &mock_image_ctx, int r) {
-    EXPECT_CALL(get_mock_rados_client(), blacklist_add(_, _))
+    EXPECT_CALL(*get_mock_io_ctx(mock_image_ctx.md_ctx).get_mock_rados_client(),
+                blacklist_add(_, _))
                   .WillOnce(Return(r));
   }
 
diff --git a/src/test/librbd/test_mock_fixture.cc b/src/test/librbd/test_mock_fixture.cc
index 97335d49b30..fba57b739cc 100644
--- a/src/test/librbd/test_mock_fixture.cc
+++ b/src/test/librbd/test_mock_fixture.cc
@@ -4,7 +4,7 @@
 #include "test/librbd/test_mock_fixture.h"
 #include "test/librbd/mock/MockImageCtx.h"
 #include "test/librados_test_stub/LibradosTestStub.h"
-#include "test/librados_test_stub/MockTestMemRadosClient.h"
+#include "test/librados_test_stub/MockTestMemCluster.h"
 
 // template definitions
 #include "librbd/AsyncRequest.cc"
@@ -21,37 +21,32 @@ using ::testing::Return;
 using ::testing::StrEq;
 using ::testing::WithArg;
 
-TestMockFixture::TestRadosClientPtr TestMockFixture::s_test_rados_client;
-::testing::NiceMock<librados::MockTestMemRadosClient> *
-  TestMockFixture::s_mock_rados_client = NULL;
+TestMockFixture::TestClusterRef TestMockFixture::s_test_cluster;
 
 void TestMockFixture::SetUpTestCase() {
-  s_test_rados_client = librados_test_stub::get_rados_client();
+  s_test_cluster = librados_test_stub::get_cluster();
 
-  // use a mock version of the in-memory rados client
-  s_mock_rados_client = new ::testing::NiceMock<librados::MockTestMemRadosClient>(
-      s_test_rados_client->cct());
-  librados_test_stub::set_rados_client(TestRadosClientPtr(s_mock_rados_client));
+  // use a mock version of the in-memory cluster
+  librados_test_stub::set_cluster(boost::shared_ptr<librados::TestCluster>(
+    new librados::MockTestMemCluster()));
   TestFixture::SetUpTestCase();
 }
 
 void TestMockFixture::TearDownTestCase() {
   TestFixture::TearDownTestCase();
-  librados_test_stub::set_rados_client(s_test_rados_client);
-  s_test_rados_client->put();
-  s_test_rados_client.reset();
-}
-
-void TestMockFixture::SetUp() {
-  TestFixture::SetUp();
+  librados_test_stub::set_cluster(s_test_cluster);
 }
 
 void TestMockFixture::TearDown() {
-  TestFixture::TearDown();
-
   // Mock rados client lives across tests -- reset it to initial state
-  ::testing::Mock::VerifyAndClear(s_mock_rados_client);
-  s_mock_rados_client->default_to_dispatch();
+  librados::MockTestMemRadosClient *mock_rados_client =
+    get_mock_io_ctx(m_ioctx).get_mock_rados_client();
+  ASSERT_TRUE(mock_rados_client != nullptr);
+
+  ::testing::Mock::VerifyAndClear(mock_rados_client);
+  mock_rados_client->default_to_dispatch();
+
+  TestFixture::TearDown();
 }
 
 void TestMockFixture::expect_unlock_exclusive_lock(librbd::ImageCtx &ictx) {
diff --git a/src/test/librbd/test_mock_fixture.h b/src/test/librbd/test_mock_fixture.h
index 6146a368610..c0bbb3d2536 100644
--- a/src/test/librbd/test_mock_fixture.h
+++ b/src/test/librbd/test_mock_fixture.h
@@ -12,7 +12,8 @@
 #include <gmock/gmock.h>
 
 namespace librados {
-class TestRadosClient;
+class TestCluster;
+class MockTestMemCluster;
 class MockTestMemIoCtxImpl;
 class MockTestMemRadosClient;
 }
@@ -59,18 +60,13 @@ MATCHER_P(ContentsEqual, bl, "") {
 
 class TestMockFixture : public TestFixture {
 public:
-  typedef boost::shared_ptr<librados::TestRadosClient> TestRadosClientPtr;
+  typedef boost::shared_ptr<librados::TestCluster> TestClusterRef;
 
   static void SetUpTestCase();
   static void TearDownTestCase();
 
-  void SetUp() override;
   void TearDown() override;
 
-  ::testing::NiceMock<librados::MockTestMemRadosClient> &get_mock_rados_client() {
-    return *s_mock_rados_client;
-  }
-
   void expect_op_work_queue(librbd::MockImageCtx &mock_image_ctx);
   void expect_unlock_exclusive_lock(librbd::ImageCtx &ictx);
 
@@ -90,8 +86,7 @@ public:
   void expect_commit_op_event(librbd::MockImageCtx &mock_image_ctx, int r);
 
 private:
-  static TestRadosClientPtr s_test_rados_client;
-  static ::testing::NiceMock<librados::MockTestMemRadosClient> *s_mock_rados_client;
+  static TestClusterRef s_test_cluster;
 };
 
 #endif // CEPH_TEST_LIBRBD_TEST_MOCK_FIXTURE_H