1 files changed, 20 insertions, 34 deletions

diff --git a/cpp/tests/ReferenceTest.cpp b/cpp/tests/ReferenceTest.cpp
index c435873585..753f68ee75 100644
--- a/cpp/tests/ReferenceTest.cpp
+++ b/cpp/tests/ReferenceTest.cpp
@@ -39,39 +39,28 @@ class ReferenceTest : public CppUnit::TestCase
     CPPUNIT_TEST(testReference);
     CPPUNIT_TEST_SUITE_END();
 
-
-    struct MockCompletionHandler : public CompletionHandler {
-        std::vector<MessageMessage::shared_ptr> messages;
-        void complete(Message::shared_ptr m) {
-            MessageMessage::shared_ptr mm =
-                dynamic_pointer_cast<MessageMessage>(m);
-            CPPUNIT_ASSERT(mm);
-            messages.push_back(mm);
-        }
-    };
-
-    MockCompletionHandler handler;
     ProtocolVersion v;
     ReferenceRegistry registry;
+    Reference::shared_ptr r1;
     MessageTransferBody::shared_ptr t1, t2;
     MessageMessage::shared_ptr m1, m2;
     MessageAppendBody::shared_ptr a1, a2;
   public:
 
     ReferenceTest() :
-        registry(handler),
+        r1(registry.open("bar")),
         t1(new MessageTransferBody(v)),
         t2(new MessageTransferBody(v)),
-        m1(new MessageMessage(0, t1)),
-        m2(new MessageMessage(0, t2)),
+        m1(new MessageMessage(0, 1, t1, r1)),
+        m2(new MessageMessage(0, 2, t2, r1)),
         a1(new MessageAppendBody(v)),
         a2(new MessageAppendBody(v))
     {}
 
     void testRegistry() {
-        Reference& ref = registry.open("foo");
-        CPPUNIT_ASSERT_EQUAL(string("foo"), ref.getId());
-        CPPUNIT_ASSERT(&ref == &registry.get("foo"));
+        Reference::shared_ptr ref = registry.open("foo");
+        CPPUNIT_ASSERT_EQUAL(string("foo"), ref->getId());
+        CPPUNIT_ASSERT(ref == registry.get("foo"));
         try {
             registry.get("none");
             CPPUNIT_FAIL("Expected exception");
@@ -83,29 +72,26 @@ class ReferenceTest : public CppUnit::TestCase
     }
 
     void testReference() {
-        Reference& ref = registry.open("foo");
-        ref.addMessage(m1);
-        ref.addMessage(m2);
-        CPPUNIT_ASSERT_EQUAL(size_t(2), ref.getMessages().size());
-        ref.append(a1);
-        ref.append(a2);
-        CPPUNIT_ASSERT_EQUAL(size_t(2), ref.getAppends().size());
-        ref.close();
+        r1->addMessage(m1);
+        r1->addMessage(m2);
+        CPPUNIT_ASSERT_EQUAL(size_t(2), r1->getMessages().size());
+        r1->append(a1);
+        r1->append(a2);
+        CPPUNIT_ASSERT_EQUAL(size_t(2), r1->getAppends().size());
+        const vector<MessageMessage::shared_ptr> messages = r1->getMessages();
+        r1->close();
         try {
-            registry.open("foo");
+            registry.open("bar");
             CPPUNIT_FAIL("Expected exception");
         } catch(...) {}
 
-        vector<MessageMessage::shared_ptr>& messages = handler.messages;
-        CPPUNIT_ASSERT_EQUAL(size_t(2), messages.size());
-
         CPPUNIT_ASSERT_EQUAL(messages[0], m1);
-        CPPUNIT_ASSERT_EQUAL(messages[0]->getAppends()[0], a1);
-        CPPUNIT_ASSERT_EQUAL(messages[0]->getAppends()[1], a2);
+        CPPUNIT_ASSERT_EQUAL(messages[0]->getReference()->getAppends()[0], a1);
+        CPPUNIT_ASSERT_EQUAL(messages[0]->getReference()->getAppends()[1], a2);
 
         CPPUNIT_ASSERT_EQUAL(messages[1], m2);
-        CPPUNIT_ASSERT_EQUAL(messages[1]->getAppends()[0], a1);
-        CPPUNIT_ASSERT_EQUAL(messages[1]->getAppends()[1], a2);
+        CPPUNIT_ASSERT_EQUAL(messages[1]->getReference()->getAppends()[0], a1);
+        CPPUNIT_ASSERT_EQUAL(messages[1]->getReference()->getAppends()[1], a2);
     }