Merge "Fix for possible NPE if Bundle is stopped."
[controller.git] / opendaylight / md-sal / sal-distributed-datastore / src / test / java / org / opendaylight / controller / cluster / datastore / TransactionProxyTest.java
index 0cd029c2ffb2da4c093ede214ca037d5a5a345dc..14696f786e7e36888be3b5517c14ae4b9779340f 100644 (file)
 package org.opendaylight.controller.cluster.datastore;
 
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.fail;
+import akka.actor.ActorPath;
 import akka.actor.ActorRef;
+import akka.actor.ActorSelection;
 import akka.actor.Props;
-
+import akka.dispatch.Futures;
 import com.google.common.base.Optional;
-import com.google.common.util.concurrent.ListenableFuture;
-import com.google.common.util.concurrent.ListeningExecutorService;
-import com.google.common.util.concurrent.MoreExecutors;
-
-import junit.framework.Assert;
 
-import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
+import org.mockito.ArgumentMatcher;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+import static org.opendaylight.controller.cluster.datastore.TransactionProxy.TransactionType.READ_ONLY;
+import static org.opendaylight.controller.cluster.datastore.TransactionProxy.TransactionType.WRITE_ONLY;
+import static org.opendaylight.controller.cluster.datastore.TransactionProxy.TransactionType.READ_WRITE;
+
+import org.opendaylight.controller.cluster.datastore.TransactionProxy.TransactionType;
 import org.opendaylight.controller.cluster.datastore.exceptions.PrimaryNotFoundException;
 import org.opendaylight.controller.cluster.datastore.exceptions.TimeoutException;
 import org.opendaylight.controller.cluster.datastore.messages.CloseTransaction;
+import org.opendaylight.controller.cluster.datastore.messages.CreateTransaction;
+import org.opendaylight.controller.cluster.datastore.messages.DataExists;
+import org.opendaylight.controller.cluster.datastore.messages.DataExistsReply;
 import org.opendaylight.controller.cluster.datastore.messages.DeleteData;
 import org.opendaylight.controller.cluster.datastore.messages.MergeData;
-import org.opendaylight.controller.cluster.datastore.messages.PrimaryFound;
+import org.opendaylight.controller.cluster.datastore.messages.ReadData;
 import org.opendaylight.controller.cluster.datastore.messages.ReadDataReply;
+import org.opendaylight.controller.cluster.datastore.messages.ReadyTransaction;
 import org.opendaylight.controller.cluster.datastore.messages.ReadyTransactionReply;
 import org.opendaylight.controller.cluster.datastore.messages.WriteData;
+import org.opendaylight.controller.cluster.datastore.shardstrategy.DefaultShardStrategy;
 import org.opendaylight.controller.cluster.datastore.shardstrategy.ShardStrategyFactory;
 import org.opendaylight.controller.cluster.datastore.utils.ActorContext;
 import org.opendaylight.controller.cluster.datastore.utils.DoNothingActor;
-import org.opendaylight.controller.cluster.datastore.utils.MessageCollectorActor;
-import org.opendaylight.controller.cluster.datastore.utils.MockActorContext;
-import org.opendaylight.controller.cluster.datastore.utils.MockClusterWrapper;
 import org.opendaylight.controller.cluster.datastore.utils.MockConfiguration;
 import org.opendaylight.controller.md.cluster.datastore.model.TestModel;
+import org.opendaylight.controller.md.sal.common.api.data.ReadFailedException;
 import org.opendaylight.controller.protobuff.messages.transaction.ShardTransactionMessages.CreateTransactionReply;
 import org.opendaylight.controller.sal.core.spi.data.DOMStoreThreePhaseCommitCohort;
 import org.opendaylight.yangtools.yang.data.api.schema.NormalizedNode;
 import org.opendaylight.yangtools.yang.data.impl.schema.ImmutableNodes;
+import org.opendaylight.yangtools.yang.model.api.SchemaContext;
+
+import scala.concurrent.Future;
 import scala.concurrent.duration.FiniteDuration;
 
-import java.util.List;
-import java.util.concurrent.Executors;
+import java.util.Arrays;
+import java.util.concurrent.TimeUnit;
 
-import static junit.framework.Assert.fail;
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.anyString;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
-
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.argThat;
+import static org.mockito.Mockito.eq;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.isA;
+
+@SuppressWarnings("resource")
 public class TransactionProxyTest extends AbstractActorTest {
 
+    @SuppressWarnings("serial")
+    static class TestException extends RuntimeException {
+    }
+
+    static interface Invoker {
+        void invoke(TransactionProxy proxy) throws Exception;
+    }
+
     private final Configuration configuration = new MockConfiguration();
 
-    private final ActorContext testContext =
-        new ActorContext(getSystem(), getSystem().actorOf(Props.create(DoNothingActor.class)), new MockClusterWrapper(), configuration );
+    @Mock
+    private ActorContext mockActorContext;
+
+    private SchemaContext schemaContext;
 
-    private final ListeningExecutorService transactionExecutor =
-        MoreExecutors.listeningDecorator(Executors.newSingleThreadExecutor());
+    String memberName = "mock-member";
 
     @Before
     public void setUp(){
+        MockitoAnnotations.initMocks(this);
+
+        schemaContext = TestModel.createTestContext();
+
+        doReturn(getSystem()).when(mockActorContext).getActorSystem();
+
         ShardStrategyFactory.setConfiguration(configuration);
     }
 
-    @After
-    public void tearDown() {
-        transactionExecutor.shutdownNow();
+    private CreateTransaction eqCreateTransaction(final String memberName,
+            final TransactionType type) {
+        ArgumentMatcher<CreateTransaction> matcher = new ArgumentMatcher<CreateTransaction>() {
+            @Override
+            public boolean matches(Object argument) {
+                CreateTransaction obj = CreateTransaction.fromSerializable(argument);
+                return obj.getTransactionId().startsWith(memberName) &&
+                       obj.getTransactionType() == type.ordinal();
+            }
+        };
+
+        return argThat(matcher);
     }
 
-    @Test
-    public void testRead() throws Exception {
-        final Props props = Props.create(DoNothingActor.class);
-        final ActorRef actorRef = getSystem().actorOf(props);
+    private DataExists eqDataExists() {
+        ArgumentMatcher<DataExists> matcher = new ArgumentMatcher<DataExists>() {
+            @Override
+            public boolean matches(Object argument) {
+                DataExists obj = DataExists.fromSerializable(argument);
+                return obj.getPath().equals(TestModel.TEST_PATH);
+            }
+        };
 
-        final MockActorContext actorContext = new MockActorContext(this.getSystem());
-        actorContext.setExecuteLocalOperationResponse(createPrimaryFound(actorRef));
-        actorContext.setExecuteShardOperationResponse(createTransactionReply(actorRef));
-        actorContext.setExecuteRemoteOperationResponse("message");
+        return argThat(matcher);
+    }
 
+    private ReadData eqReadData() {
+        ArgumentMatcher<ReadData> matcher = new ArgumentMatcher<ReadData>() {
+            @Override
+            public boolean matches(Object argument) {
+                ReadData obj = ReadData.fromSerializable(argument);
+                return obj.getPath().equals(TestModel.TEST_PATH);
+            }
+        };
 
-        TransactionProxy transactionProxy =
-            new TransactionProxy(actorContext,
-                TransactionProxy.TransactionType.READ_ONLY, transactionExecutor, TestModel.createTestContext());
+        return argThat(matcher);
+    }
 
+    private WriteData eqWriteData(final NormalizedNode<?, ?> nodeToWrite) {
+        ArgumentMatcher<WriteData> matcher = new ArgumentMatcher<WriteData>() {
+            @Override
+            public boolean matches(Object argument) {
+                WriteData obj = WriteData.fromSerializable(argument, schemaContext);
+                return obj.getPath().equals(TestModel.TEST_PATH) &&
+                       obj.getData().equals(nodeToWrite);
+            }
+        };
+
+        return argThat(matcher);
+    }
 
-        ListenableFuture<Optional<NormalizedNode<?, ?>>> read =
-            transactionProxy.read(TestModel.TEST_PATH);
+    private MergeData eqMergeData(final NormalizedNode<?, ?> nodeToWrite) {
+        ArgumentMatcher<MergeData> matcher = new ArgumentMatcher<MergeData>() {
+            @Override
+            public boolean matches(Object argument) {
+                MergeData obj = MergeData.fromSerializable(argument, schemaContext);
+                return obj.getPath().equals(TestModel.TEST_PATH) &&
+                       obj.getData().equals(nodeToWrite);
+            }
+        };
+
+        return argThat(matcher);
+    }
 
-        Optional<NormalizedNode<?, ?>> normalizedNodeOptional = read.get();
+    private DeleteData eqDeleteData() {
+        ArgumentMatcher<DeleteData> matcher = new ArgumentMatcher<DeleteData>() {
+            @Override
+            public boolean matches(Object argument) {
+                DeleteData obj = DeleteData.fromSerializable(argument);
+                return obj.getPath().equals(TestModel.TEST_PATH);
+            }
+        };
 
-        Assert.assertFalse(normalizedNodeOptional.isPresent());
+        return argThat(matcher);
+    }
 
-        actorContext.setExecuteRemoteOperationResponse(new ReadDataReply(
-            TestModel.createTestContext(),ImmutableNodes.containerNode(TestModel.TEST_QNAME)).toSerializable());
+    private Object readyTxReply(ActorPath path) {
+        return new ReadyTransactionReply(path).toSerializable();
+    }
 
-        read = transactionProxy.read(TestModel.TEST_PATH);
+    private Future<Object> readDataReply(NormalizedNode<?, ?> data) {
+        return Futures.successful(new ReadDataReply(schemaContext, data)
+                .toSerializable());
+    }
 
-        normalizedNodeOptional = read.get();
+    private Future<Object> dataExistsReply(boolean exists) {
+        return Futures.successful(new DataExistsReply(exists).toSerializable());
+    }
 
-        Assert.assertTrue(normalizedNodeOptional.isPresent());
+    private ActorSelection actorSelection(ActorRef actorRef) {
+        return getSystem().actorSelection(actorRef.path());
     }
 
-    @Test
-    public void testReadWhenANullIsReturned() throws Exception {
-        final Props props = Props.create(DoNothingActor.class);
-        final ActorRef actorRef = getSystem().actorOf(props);
+    private FiniteDuration anyDuration() {
+        return any(FiniteDuration.class);
+    }
 
-        final MockActorContext actorContext = new MockActorContext(this.getSystem());
-        actorContext.setExecuteLocalOperationResponse(createPrimaryFound(actorRef));
-        actorContext.setExecuteShardOperationResponse(createTransactionReply(actorRef));
-        actorContext.setExecuteRemoteOperationResponse("message");
+    private CreateTransactionReply createTransactionReply(ActorRef actorRef){
+        return CreateTransactionReply.newBuilder()
+            .setTransactionActorPath(actorRef.path().toString())
+            .setTransactionId("txn-1").build();
+    }
 
-        TransactionProxy transactionProxy =
-            new TransactionProxy(actorContext,
-                TransactionProxy.TransactionType.READ_ONLY, transactionExecutor, TestModel.createTestContext());
+    private ActorRef setupActorContextWithInitialCreateTransaction(TransactionType type) {
+        ActorRef actorRef = getSystem().actorOf(Props.create(DoNothingActor.class));
+        doReturn(getSystem().actorSelection(actorRef.path())).
+                when(mockActorContext).actorSelection(actorRef.path().toString());
+        doReturn(memberName).when(mockActorContext).getCurrentMemberName();
+        doReturn(createTransactionReply(actorRef)).when(mockActorContext).
+                executeShardOperation(eq(DefaultShardStrategy.DEFAULT_SHARD),
+                        eqCreateTransaction(memberName, type), anyDuration());
+        doReturn(actorRef.path().toString()).when(mockActorContext).resolvePath(
+                anyString(), eq(actorRef.path().toString()));
+        doReturn(actorRef.path()).when(mockActorContext).actorFor(actorRef.path().toString());
+
+        return actorRef;
+    }
 
+    @Test
+    public void testRead() throws Exception {
+        ActorRef actorRef = setupActorContextWithInitialCreateTransaction(READ_ONLY);
 
-        ListenableFuture<Optional<NormalizedNode<?, ?>>> read =
-            transactionProxy.read(TestModel.TEST_PATH);
+        TransactionProxy transactionProxy = new TransactionProxy(mockActorContext,
+                READ_ONLY, schemaContext);
 
-        Optional<NormalizedNode<?, ?>> normalizedNodeOptional = read.get();
+        doReturn(readDataReply(null)).when(mockActorContext).executeRemoteOperationAsync(
+                eq(actorSelection(actorRef)), eqReadData(), anyDuration());
 
-        Assert.assertFalse(normalizedNodeOptional.isPresent());
+        Optional<NormalizedNode<?, ?>> readOptional = transactionProxy.read(
+                TestModel.TEST_PATH).get(5, TimeUnit.SECONDS);
 
-        actorContext.setExecuteRemoteOperationResponse(new ReadDataReply(
-           TestModel.createTestContext(), null).toSerializable());
+        assertEquals("NormalizedNode isPresent", false, readOptional.isPresent());
 
-        read = transactionProxy.read(TestModel.TEST_PATH);
+        NormalizedNode<?, ?> expectedNode = ImmutableNodes.containerNode(TestModel.TEST_QNAME);
 
-        normalizedNodeOptional = read.get();
+        doReturn(readDataReply(expectedNode)).when(mockActorContext).executeRemoteOperationAsync(
+                eq(actorSelection(actorRef)), eqReadData(), anyDuration());
 
-        Assert.assertFalse(normalizedNodeOptional.isPresent());
+        readOptional = transactionProxy.read(TestModel.TEST_PATH).get(5, TimeUnit.SECONDS);
+
+        assertEquals("NormalizedNode isPresent", true, readOptional.isPresent());
+
+        assertEquals("Response NormalizedNode", expectedNode, readOptional.get());
     }
 
-    @Test
-    public void testReadWhenAPrimaryNotFoundExceptionIsThrown() throws Exception {
-        final ActorContext actorContext = mock(ActorContext.class);
+    @Test(expected = ReadFailedException.class)
+    public void testReadWhenAnInvalidMessageIsSentInReply() throws Exception {
+        setupActorContextWithInitialCreateTransaction(READ_ONLY);
 
-        when(actorContext.executeShardOperation(anyString(), any(), any(
-            FiniteDuration.class))).thenThrow(new PrimaryNotFoundException("test"));
+        doReturn(Futures.successful(new Object())).when(mockActorContext).
+                executeRemoteOperationAsync(any(ActorSelection.class), any(), anyDuration());
 
-        TransactionProxy transactionProxy =
-            new TransactionProxy(actorContext,
-                TransactionProxy.TransactionType.READ_ONLY, transactionExecutor, TestModel.createTestContext());
+        TransactionProxy transactionProxy = new TransactionProxy(mockActorContext,
+                READ_ONLY, schemaContext);
+
+        transactionProxy.read(TestModel.TEST_PATH).checkedGet(5, TimeUnit.SECONDS);
+    }
 
+    @Test(expected = TestException.class)
+    public void testReadWithAsyncRemoteOperatonFailure() throws Throwable {
+        setupActorContextWithInitialCreateTransaction(READ_ONLY);
 
-        ListenableFuture<Optional<NormalizedNode<?, ?>>> read =
-            transactionProxy.read(TestModel.TEST_PATH);
+        doThrow(new TestException()).when(mockActorContext).
+                executeRemoteOperationAsync(any(ActorSelection.class), any(), anyDuration());
 
-        Assert.assertFalse(read.get().isPresent());
+        TransactionProxy transactionProxy = new TransactionProxy(mockActorContext,
+                READ_ONLY, schemaContext);
 
+        try {
+            transactionProxy.read(TestModel.TEST_PATH).checkedGet(5, TimeUnit.SECONDS);
+            fail("Expected ReadFailedException");
+        } catch(ReadFailedException e) {
+            // Expected - throw cause - expects TestException.
+            throw e.getCause();
+        }
     }
 
+    private void testExceptionOnInitialCreateTransaction(Exception exToThrow, Invoker invoker)
+            throws Throwable {
 
-    @Test
-    public void testReadWhenATimeoutExceptionIsThrown() throws Exception {
-        final ActorContext actorContext = mock(ActorContext.class);
+        doThrow(exToThrow).when(mockActorContext).executeShardOperation(
+                anyString(), any(), anyDuration());
 
-        when(actorContext.executeShardOperation(anyString(), any(), any(
-            FiniteDuration.class))).thenThrow(new TimeoutException("test", new Exception("reason")));
+        TransactionProxy transactionProxy = new TransactionProxy(mockActorContext,
+                READ_ONLY, schemaContext);
 
-        TransactionProxy transactionProxy =
-            new TransactionProxy(actorContext,
-                TransactionProxy.TransactionType.READ_ONLY, transactionExecutor, TestModel.createTestContext());
+        try {
+            invoker.invoke(transactionProxy);
+            fail("Expected ReadFailedException");
+        } catch(ReadFailedException e) {
+            // Expected - throw cause - expects TestException.
+            throw e.getCause();
+        }
+    }
 
+    private void testReadWithExceptionOnInitialCreateTransaction(Exception exToThrow) throws Throwable {
+        testExceptionOnInitialCreateTransaction(exToThrow, new Invoker() {
+            @Override
+            public void invoke(TransactionProxy proxy) throws Exception {
+                proxy.read(TestModel.TEST_PATH).checkedGet(5, TimeUnit.SECONDS);
+            }
+        });
+    }
 
-        ListenableFuture<Optional<NormalizedNode<?, ?>>> read =
-            transactionProxy.read(TestModel.TEST_PATH);
+    @Test(expected = PrimaryNotFoundException.class)
+    public void testReadWhenAPrimaryNotFoundExceptionIsThrown() throws Throwable {
+        testReadWithExceptionOnInitialCreateTransaction(new PrimaryNotFoundException("test"));
+    }
 
-        Assert.assertFalse(read.get().isPresent());
+    @Test(expected = TimeoutException.class)
+    public void testReadWhenATimeoutExceptionIsThrown() throws Throwable {
+        testReadWithExceptionOnInitialCreateTransaction(new TimeoutException("test",
+                new Exception("reason")));
+    }
 
+    @Test(expected = TestException.class)
+    public void testReadWhenAnyOtherExceptionIsThrown() throws Throwable {
+        testReadWithExceptionOnInitialCreateTransaction(new TestException());
     }
 
     @Test
-    public void testReadWhenAAnyOtherExceptionIsThrown() throws Exception {
-        final ActorContext actorContext = mock(ActorContext.class);
+    public void testExists() throws Exception {
+        ActorRef actorRef = setupActorContextWithInitialCreateTransaction(READ_ONLY);
 
-        when(actorContext.executeShardOperation(anyString(), any(), any(
-            FiniteDuration.class))).thenThrow(new NullPointerException());
+        TransactionProxy transactionProxy = new TransactionProxy(mockActorContext,
+                READ_ONLY, schemaContext);
 
-        TransactionProxy transactionProxy =
-            new TransactionProxy(actorContext,
-                TransactionProxy.TransactionType.READ_ONLY, transactionExecutor, TestModel.createTestContext());
+        doReturn(dataExistsReply(false)).when(mockActorContext).executeRemoteOperationAsync(
+                eq(actorSelection(actorRef)), eqDataExists(), anyDuration());
 
+        Boolean exists = transactionProxy.exists(TestModel.TEST_PATH).checkedGet();
 
-        try {
-            ListenableFuture<Optional<NormalizedNode<?, ?>>> read =
-                transactionProxy.read(TestModel.TEST_PATH);
-            fail("A null pointer exception was expected");
-        } catch(NullPointerException e){
-
-        }
-    }
-
+        assertEquals("Exists response", false, exists);
 
+        doReturn(dataExistsReply(true)).when(mockActorContext).executeRemoteOperationAsync(
+                eq(actorSelection(actorRef)), eqDataExists(), anyDuration());
 
-    @Test
-    public void testWrite() throws Exception {
-        final Props props = Props.create(MessageCollectorActor.class);
-        final ActorRef actorRef = getSystem().actorOf(props);
+        exists = transactionProxy.exists(TestModel.TEST_PATH).checkedGet();
 
-        final MockActorContext actorContext = new MockActorContext(this.getSystem());
-        actorContext.setExecuteLocalOperationResponse(createPrimaryFound(actorRef));
-        actorContext.setExecuteShardOperationResponse(createTransactionReply(actorRef));
-        actorContext.setExecuteRemoteOperationResponse("message");
+        assertEquals("Exists response", true, exists);
+    }
 
-        TransactionProxy transactionProxy =
-            new TransactionProxy(actorContext,
-                TransactionProxy.TransactionType.READ_ONLY, transactionExecutor, TestModel.createTestContext());
+    @Test(expected = PrimaryNotFoundException.class)
+    public void testExistsWhenAPrimaryNotFoundExceptionIsThrown() throws Throwable {
+        testExceptionOnInitialCreateTransaction(new PrimaryNotFoundException("test"), new Invoker() {
+            @Override
+            public void invoke(TransactionProxy proxy) throws Exception {
+                proxy.exists(TestModel.TEST_PATH).checkedGet(5, TimeUnit.SECONDS);
+            }
+        });
+    }
 
-        transactionProxy.write(TestModel.TEST_PATH,
-            ImmutableNodes.containerNode(TestModel.NAME_QNAME));
+    @Test(expected = ReadFailedException.class)
+    public void testExistsWhenAnInvalidMessageIsSentInReply() throws Exception {
+        setupActorContextWithInitialCreateTransaction(READ_ONLY);
 
-        Object messages = testContext
-            .executeLocalOperation(actorRef, "messages",
-                ActorContext.ASK_DURATION);
+        doReturn(Futures.successful(new Object())).when(mockActorContext).
+                executeRemoteOperationAsync(any(ActorSelection.class), any(), anyDuration());
 
-        Assert.assertNotNull(messages);
+        TransactionProxy transactionProxy = new TransactionProxy(mockActorContext,
+                READ_ONLY, schemaContext);
 
-        Assert.assertTrue(messages instanceof List);
+        transactionProxy.exists(TestModel.TEST_PATH).checkedGet(5, TimeUnit.SECONDS);
+    }
 
-        List<Object> listMessages = (List<Object>) messages;
+    @Test(expected = TestException.class)
+    public void testExistsWithAsyncRemoteOperatonFailure() throws Throwable {
+        setupActorContextWithInitialCreateTransaction(READ_ONLY);
 
-        Assert.assertEquals(1, listMessages.size());
+        doThrow(new TestException()).when(mockActorContext).
+                executeRemoteOperationAsync(any(ActorSelection.class), any(), anyDuration());
 
-        Assert.assertEquals(WriteData.SERIALIZABLE_CLASS, listMessages.get(0).getClass());
-    }
+        TransactionProxy transactionProxy = new TransactionProxy(mockActorContext,
+                READ_ONLY, schemaContext);
 
-    private Object createPrimaryFound(ActorRef actorRef) {
-        return new PrimaryFound(actorRef.path().toString()).toSerializable();
+        try {
+            transactionProxy.exists(TestModel.TEST_PATH).checkedGet(5, TimeUnit.SECONDS);
+            fail("Expected ReadFailedException");
+        } catch(ReadFailedException e) {
+            // Expected - throw cause - expects TestException.
+            throw e.getCause();
+        }
     }
 
     @Test
-    public void testMerge() throws Exception {
-        final Props props = Props.create(MessageCollectorActor.class);
-        final ActorRef actorRef = getSystem().actorOf(props);
+    public void testWrite() throws Exception {
+        ActorRef actorRef = setupActorContextWithInitialCreateTransaction(WRITE_ONLY);
 
-        final MockActorContext actorContext = new MockActorContext(this.getSystem());
-        actorContext.setExecuteLocalOperationResponse(createPrimaryFound(actorRef));
-        actorContext.setExecuteShardOperationResponse(createTransactionReply(actorRef));
-        actorContext.setExecuteRemoteOperationResponse("message");
+        TransactionProxy transactionProxy = new TransactionProxy(mockActorContext,
+                WRITE_ONLY, schemaContext);
 
-        TransactionProxy transactionProxy =
-            new TransactionProxy(actorContext,
-                TransactionProxy.TransactionType.READ_ONLY, transactionExecutor, TestModel.createTestContext());
+        NormalizedNode<?, ?> nodeToWrite = ImmutableNodes.containerNode(TestModel.TEST_QNAME);
 
-        transactionProxy.merge(TestModel.TEST_PATH,
-            ImmutableNodes.containerNode(TestModel.NAME_QNAME));
+        transactionProxy.write(TestModel.TEST_PATH, nodeToWrite);
 
-        Object messages = testContext
-            .executeLocalOperation(actorRef, "messages",
-                ActorContext.ASK_DURATION);
+        verify(mockActorContext).sendRemoteOperationAsync(
+                eq(actorSelection(actorRef)), eqWriteData(nodeToWrite));
+    }
 
-        Assert.assertNotNull(messages);
+    @Test
+    public void testMerge() throws Exception {
+        ActorRef actorRef = setupActorContextWithInitialCreateTransaction(WRITE_ONLY);
 
-        Assert.assertTrue(messages instanceof List);
+        TransactionProxy transactionProxy = new TransactionProxy(mockActorContext,
+                WRITE_ONLY, schemaContext);
 
-        List<Object> listMessages = (List<Object>) messages;
+        NormalizedNode<?, ?> nodeToWrite = ImmutableNodes.containerNode(TestModel.TEST_QNAME);
 
-        Assert.assertEquals(1, listMessages.size());
+        transactionProxy.merge(TestModel.TEST_PATH, nodeToWrite);
 
-        Assert.assertEquals(MergeData.SERIALIZABLE_CLASS, listMessages.get(0).getClass());
+        verify(mockActorContext).sendRemoteOperationAsync(
+                eq(actorSelection(actorRef)), eqMergeData(nodeToWrite));
     }
 
     @Test
     public void testDelete() throws Exception {
-        final Props props = Props.create(MessageCollectorActor.class);
-        final ActorRef actorRef = getSystem().actorOf(props);
-
-        final MockActorContext actorContext = new MockActorContext(this.getSystem());
-        actorContext.setExecuteLocalOperationResponse(createPrimaryFound(actorRef));
-        actorContext.setExecuteShardOperationResponse(createTransactionReply(actorRef));
-        actorContext.setExecuteRemoteOperationResponse("message");
+        ActorRef actorRef = setupActorContextWithInitialCreateTransaction(WRITE_ONLY);
 
-        TransactionProxy transactionProxy =
-            new TransactionProxy(actorContext,
-                TransactionProxy.TransactionType.READ_ONLY, transactionExecutor, TestModel.createTestContext());
+        TransactionProxy transactionProxy = new TransactionProxy(mockActorContext,
+                WRITE_ONLY, schemaContext);
 
         transactionProxy.delete(TestModel.TEST_PATH);
 
-        Object messages = testContext
-            .executeLocalOperation(actorRef, "messages",
-                ActorContext.ASK_DURATION);
-
-        Assert.assertNotNull(messages);
-
-        Assert.assertTrue(messages instanceof List);
-
-        List<Object> listMessages = (List<Object>) messages;
-
-        Assert.assertEquals(1, listMessages.size());
-
-        Assert.assertEquals(DeleteData.SERIALIZABLE_CLASS, listMessages.get(0).getClass());
+        verify(mockActorContext).sendRemoteOperationAsync(
+                eq(actorSelection(actorRef)), eqDeleteData());
     }
 
+    @SuppressWarnings("unchecked")
     @Test
     public void testReady() throws Exception {
-        final Props props = Props.create(DoNothingActor.class);
-        final ActorRef doNothingActorRef = getSystem().actorOf(props);
+        ActorRef actorRef = setupActorContextWithInitialCreateTransaction(READ_WRITE);
 
-        final MockActorContext actorContext = new MockActorContext(this.getSystem());
-        actorContext.setExecuteLocalOperationResponse(createPrimaryFound(doNothingActorRef));
-        actorContext.setExecuteShardOperationResponse(createTransactionReply(doNothingActorRef));
-        actorContext.setExecuteRemoteOperationResponse(new ReadyTransactionReply(doNothingActorRef.path()).toSerializable());
+        doReturn(readDataReply(null)).when(mockActorContext).executeRemoteOperationAsync(
+                eq(actorSelection(actorRef)), eqReadData(), anyDuration());
 
-        TransactionProxy transactionProxy =
-            new TransactionProxy(actorContext,
-                TransactionProxy.TransactionType.READ_ONLY, transactionExecutor, TestModel.createTestContext());
+        doReturn(readyTxReply(actorRef.path())).when(mockActorContext).executeRemoteOperation(
+                eq(actorSelection(actorRef)), isA(ReadyTransaction.SERIALIZABLE_CLASS), anyDuration());
 
+        TransactionProxy transactionProxy = new TransactionProxy(mockActorContext,
+                READ_WRITE, schemaContext);
 
         transactionProxy.read(TestModel.TEST_PATH);
 
         DOMStoreThreePhaseCommitCohort ready = transactionProxy.ready();
 
-        Assert.assertTrue(ready instanceof ThreePhaseCommitCohortProxy);
+        assertTrue(ready instanceof ThreePhaseCommitCohortProxy);
 
         ThreePhaseCommitCohortProxy proxy = (ThreePhaseCommitCohortProxy) ready;
 
-        Assert.assertTrue("No cohort paths returned", proxy.getCohortPaths().size() > 0);
-
+        assertEquals("getCohortPaths", Arrays.asList(actorRef.path()), proxy.getCohortPaths());
     }
 
     @Test
-    public void testGetIdentifier(){
-        final Props props = Props.create(DoNothingActor.class);
-        final ActorRef doNothingActorRef = getSystem().actorOf(props);
-
-        final MockActorContext actorContext = new MockActorContext(this.getSystem());
-        actorContext.setExecuteShardOperationResponse( createTransactionReply(doNothingActorRef) );
-
-        TransactionProxy transactionProxy =
-            new TransactionProxy(actorContext,
-                TransactionProxy.TransactionType.READ_ONLY, transactionExecutor, TestModel.createTestContext());
-
-        Assert.assertNotNull(transactionProxy.getIdentifier());
+    public void testGetIdentifier() {
+        setupActorContextWithInitialCreateTransaction(READ_ONLY);
+        TransactionProxy transactionProxy = new TransactionProxy(mockActorContext,
+                TransactionProxy.TransactionType.READ_ONLY, schemaContext);
+
+        Object id = transactionProxy.getIdentifier();
+        assertNotNull("getIdentifier returned null", id);
+        assertTrue("Invalid identifier: " + id, id.toString().startsWith(memberName));
     }
 
+    @SuppressWarnings("unchecked")
     @Test
-    public void testClose(){
-        final Props props = Props.create(MessageCollectorActor.class);
-        final ActorRef actorRef = getSystem().actorOf(props);
+    public void testClose() throws Exception{
+        ActorRef actorRef = setupActorContextWithInitialCreateTransaction(READ_WRITE);
 
-        final MockActorContext actorContext = new MockActorContext(this.getSystem());
-        actorContext.setExecuteLocalOperationResponse(createPrimaryFound(actorRef));
-        actorContext.setExecuteShardOperationResponse(createTransactionReply(actorRef));
-        actorContext.setExecuteRemoteOperationResponse("message");
+        doReturn(readDataReply(null)).when(mockActorContext).executeRemoteOperationAsync(
+                eq(actorSelection(actorRef)), eqReadData(), anyDuration());
 
-        TransactionProxy transactionProxy =
-            new TransactionProxy(actorContext,
-                TransactionProxy.TransactionType.READ_ONLY, transactionExecutor, TestModel.createTestContext());
+        TransactionProxy transactionProxy = new TransactionProxy(mockActorContext,
+                READ_WRITE, schemaContext);
 
         transactionProxy.read(TestModel.TEST_PATH);
 
         transactionProxy.close();
 
-        Object messages = testContext
-            .executeLocalOperation(actorRef, "messages",
-                ActorContext.ASK_DURATION);
-
-        Assert.assertNotNull(messages);
-
-        Assert.assertTrue(messages instanceof List);
-
-        List<Object> listMessages = (List<Object>) messages;
-
-        Assert.assertEquals(1, listMessages.size());
-
-        Assert.assertTrue(listMessages.get(0).getClass().equals(CloseTransaction.SERIALIZABLE_CLASS));
-    }
-
-    private CreateTransactionReply createTransactionReply(ActorRef actorRef){
-        return CreateTransactionReply.newBuilder()
-            .setTransactionActorPath(actorRef.path().toString())
-            .setTransactionId("txn-1")
-            .build();
+        verify(mockActorContext).sendRemoteOperationAsync(
+                eq(actorSelection(actorRef)), isA(CloseTransaction.SERIALIZABLE_CLASS));
     }
 }