+ long start = System.nanoTime();
+
+ operation.run(transactionProxy);
+
+ long end = System.nanoTime();
+
+ long expected = TimeUnit.MILLISECONDS.toNanos(mockActorContext.getDatastoreContext().getOperationTimeoutInMillis());
+ Assert.assertTrue(String.format("Expected elapsed time: %s. Actual: %s",
+ expected, (end-start)), (end - start) <= expected);
+ }
+
+ private void completeOperationLocal(TransactionProxyOperation operation, Optional<DataTree> dataTreeOptional){
+ ActorSystem actorSystem = getSystem();
+ ActorRef shardActorRef = actorSystem.actorOf(Props.create(DoNothingActor.class));
+
+ doReturn(actorSystem.actorSelection(shardActorRef.path())).
+ when(mockActorContext).actorSelection(shardActorRef.path().toString());
+
+ doReturn(Futures.successful(newPrimaryShardInfo(shardActorRef, dataTreeOptional))).
+ when(mockActorContext).findPrimaryShardAsync(eq(DefaultShardStrategy.DEFAULT_SHARD));
+
+ TransactionProxy transactionProxy = new TransactionProxy(mockComponentFactory, READ_WRITE);
+
+ long start = System.nanoTime();
+
+ operation.run(transactionProxy);
+
+ long end = System.nanoTime();
+
+ long expected = TimeUnit.MILLISECONDS.toNanos(mockActorContext.getDatastoreContext().getOperationTimeoutInMillis());
+ Assert.assertTrue(String.format("Expected elapsed time: %s. Actual: %s",
+ expected, (end-start)), (end - start) <= expected);
+ }
+
+ private static Optional<DataTree> createDataTree(){
+ DataTree dataTree = mock(DataTree.class);
+ Optional<DataTree> dataTreeOptional = Optional.of(dataTree);
+ DataTreeSnapshot dataTreeSnapshot = mock(DataTreeSnapshot.class);
+ DataTreeModification dataTreeModification = mock(DataTreeModification.class);
+
+ doReturn(dataTreeSnapshot).when(dataTree).takeSnapshot();
+ doReturn(dataTreeModification).when(dataTreeSnapshot).newModification();
+
+ return dataTreeOptional;
+ }
+
+ private static Optional<DataTree> createDataTree(NormalizedNode readResponse){
+ DataTree dataTree = mock(DataTree.class);
+ Optional<DataTree> dataTreeOptional = Optional.of(dataTree);
+ DataTreeSnapshot dataTreeSnapshot = mock(DataTreeSnapshot.class);
+ DataTreeModification dataTreeModification = mock(DataTreeModification.class);
+
+ doReturn(dataTreeSnapshot).when(dataTree).takeSnapshot();
+ doReturn(dataTreeModification).when(dataTreeSnapshot).newModification();
+ doReturn(Optional.of(readResponse)).when(dataTreeModification).readNode(any(YangInstanceIdentifier.class));
+
+ return dataTreeOptional;
+ }
+
+
+ @Test
+ public void testWriteCompletionForLocalShard(){
+ completeOperationLocal(new TransactionProxyOperation() {
+ @Override
+ public void run(TransactionProxy transactionProxy) {
+ NormalizedNode<?, ?> nodeToWrite = ImmutableNodes.containerNode(TestModel.TEST_QNAME);
+
+ transactionProxy.write(TestModel.TEST_PATH, nodeToWrite);
+
+ transactionProxy.write(TestModel.TEST_PATH, nodeToWrite);
+
+ }
+ }, createDataTree());
+ }
+
+ @Test
+ public void testWriteThrottlingWhenShardFound(){
+ throttleOperation(new TransactionProxyOperation() {
+ @Override
+ public void run(TransactionProxy transactionProxy) {
+ NormalizedNode<?, ?> nodeToWrite = ImmutableNodes.containerNode(TestModel.TEST_QNAME);
+
+ expectIncompleteBatchedModifications();
+
+ transactionProxy.write(TestModel.TEST_PATH, nodeToWrite);
+
+ transactionProxy.write(TestModel.TEST_PATH, nodeToWrite);
+ }
+ });
+ }
+
+ @Test
+ public void testWriteThrottlingWhenShardNotFound(){
+ // Confirm that there is no throttling when the Shard is not found
+ completeOperation(new TransactionProxyOperation() {
+ @Override
+ public void run(TransactionProxy transactionProxy) {
+ NormalizedNode<?, ?> nodeToWrite = ImmutableNodes.containerNode(TestModel.TEST_QNAME);
+
+ expectBatchedModifications(2);
+
+ transactionProxy.write(TestModel.TEST_PATH, nodeToWrite);
+
+ transactionProxy.write(TestModel.TEST_PATH, nodeToWrite);
+ }
+ }, false);
+
+ }
+
+
+ @Test
+ public void testWriteCompletion(){
+ completeOperation(new TransactionProxyOperation() {
+ @Override
+ public void run(TransactionProxy transactionProxy) {
+ NormalizedNode<?, ?> nodeToWrite = ImmutableNodes.containerNode(TestModel.TEST_QNAME);
+
+ expectBatchedModifications(2);
+
+ transactionProxy.write(TestModel.TEST_PATH, nodeToWrite);
+
+ transactionProxy.write(TestModel.TEST_PATH, nodeToWrite);
+ }
+ });
+ }
+
+ @Test
+ public void testMergeThrottlingWhenShardFound(){
+ throttleOperation(new TransactionProxyOperation() {
+ @Override
+ public void run(TransactionProxy transactionProxy) {
+ NormalizedNode<?, ?> nodeToMerge = ImmutableNodes.containerNode(TestModel.TEST_QNAME);
+
+ expectIncompleteBatchedModifications();
+
+ transactionProxy.merge(TestModel.TEST_PATH, nodeToMerge);
+
+ transactionProxy.merge(TestModel.TEST_PATH, nodeToMerge);
+ }
+ });
+ }
+
+ @Test
+ public void testMergeThrottlingWhenShardNotFound(){
+ completeOperation(new TransactionProxyOperation() {
+ @Override
+ public void run(TransactionProxy transactionProxy) {
+ NormalizedNode<?, ?> nodeToMerge = ImmutableNodes.containerNode(TestModel.TEST_QNAME);
+
+ expectBatchedModifications(2);
+
+ transactionProxy.merge(TestModel.TEST_PATH, nodeToMerge);
+
+ transactionProxy.merge(TestModel.TEST_PATH, nodeToMerge);
+ }
+ }, false);
+ }
+
+ @Test
+ public void testMergeCompletion(){
+ completeOperation(new TransactionProxyOperation() {
+ @Override
+ public void run(TransactionProxy transactionProxy) {
+ NormalizedNode<?, ?> nodeToMerge = ImmutableNodes.containerNode(TestModel.TEST_QNAME);
+
+ expectBatchedModifications(2);
+
+ transactionProxy.merge(TestModel.TEST_PATH, nodeToMerge);
+
+ transactionProxy.merge(TestModel.TEST_PATH, nodeToMerge);
+ }
+ });
+
+ }
+
+ @Test
+ public void testMergeCompletionForLocalShard(){
+ completeOperationLocal(new TransactionProxyOperation() {
+ @Override
+ public void run(TransactionProxy transactionProxy) {
+ NormalizedNode<?, ?> nodeToWrite = ImmutableNodes.containerNode(TestModel.TEST_QNAME);
+
+ transactionProxy.merge(TestModel.TEST_PATH, nodeToWrite);
+
+ transactionProxy.merge(TestModel.TEST_PATH, nodeToWrite);
+
+ }
+ }, createDataTree());
+ }
+
+
+ @Test
+ public void testDeleteThrottlingWhenShardFound(){
+
+ throttleOperation(new TransactionProxyOperation() {
+ @Override
+ public void run(TransactionProxy transactionProxy) {
+ expectIncompleteBatchedModifications();
+
+ transactionProxy.delete(TestModel.TEST_PATH);
+
+ transactionProxy.delete(TestModel.TEST_PATH);
+ }
+ });
+ }
+
+
+ @Test
+ public void testDeleteThrottlingWhenShardNotFound(){
+
+ completeOperation(new TransactionProxyOperation() {
+ @Override
+ public void run(TransactionProxy transactionProxy) {
+ expectBatchedModifications(2);
+
+ transactionProxy.delete(TestModel.TEST_PATH);
+
+ transactionProxy.delete(TestModel.TEST_PATH);
+ }
+ }, false);
+ }
+
+ @Test
+ public void testDeleteCompletionForLocalShard(){
+ completeOperationLocal(new TransactionProxyOperation() {
+ @Override
+ public void run(TransactionProxy transactionProxy) {
+
+ transactionProxy.delete(TestModel.TEST_PATH);
+
+ transactionProxy.delete(TestModel.TEST_PATH);
+ }
+ }, createDataTree());
+
+ }
+
+ @Test
+ public void testDeleteCompletion(){
+ completeOperation(new TransactionProxyOperation() {
+ @Override
+ public void run(TransactionProxy transactionProxy) {
+ expectBatchedModifications(2);
+
+ transactionProxy.delete(TestModel.TEST_PATH);
+
+ transactionProxy.delete(TestModel.TEST_PATH);
+ }
+ });
+
+ }
+
+ @Test
+ public void testReadThrottlingWhenShardFound(){
+
+ throttleOperation(new TransactionProxyOperation() {
+ @Override
+ public void run(TransactionProxy transactionProxy) {
+ doReturn(incompleteFuture()).when(mockActorContext).executeOperationAsync(
+ any(ActorSelection.class), eqReadData());
+
+ transactionProxy.read(TestModel.TEST_PATH);
+
+ transactionProxy.read(TestModel.TEST_PATH);
+ }
+ });
+ }
+
+ @Test
+ public void testReadThrottlingWhenShardNotFound(){
+
+ completeOperation(new TransactionProxyOperation() {
+ @Override
+ public void run(TransactionProxy transactionProxy) {
+ doReturn(incompleteFuture()).when(mockActorContext).executeOperationAsync(
+ any(ActorSelection.class), eqReadData());
+
+ transactionProxy.read(TestModel.TEST_PATH);
+
+ transactionProxy.read(TestModel.TEST_PATH);
+ }
+ }, false);
+ }
+
+
+ @Test
+ public void testReadCompletion(){
+ completeOperation(new TransactionProxyOperation() {
+ @Override
+ public void run(TransactionProxy transactionProxy) {
+ NormalizedNode<?, ?> nodeToRead = ImmutableNodes.containerNode(TestModel.TEST_QNAME);
+
+ doReturn(readDataReply(nodeToRead)).when(mockActorContext).executeOperationAsync(
+ any(ActorSelection.class), eqReadData());
+
+ transactionProxy.read(TestModel.TEST_PATH);
+
+ transactionProxy.read(TestModel.TEST_PATH);
+ }
+ });
+
+ }
+
+ @Test
+ public void testReadCompletionForLocalShard(){
+ final NormalizedNode nodeToRead = ImmutableNodes.containerNode(TestModel.TEST_QNAME);
+ completeOperationLocal(new TransactionProxyOperation() {
+ @Override
+ public void run(TransactionProxy transactionProxy) {
+ transactionProxy.read(TestModel.TEST_PATH);
+
+ transactionProxy.read(TestModel.TEST_PATH);
+ }
+ }, createDataTree(nodeToRead));
+
+ }
+
+ @Test
+ public void testReadCompletionForLocalShardWhenExceptionOccurs(){
+ completeOperationLocal(new TransactionProxyOperation() {
+ @Override
+ public void run(TransactionProxy transactionProxy) {
+ transactionProxy.read(TestModel.TEST_PATH);
+
+ transactionProxy.read(TestModel.TEST_PATH);
+ }
+ }, createDataTree());
+
+ }
+
+ @Test
+ public void testExistsThrottlingWhenShardFound(){
+
+ throttleOperation(new TransactionProxyOperation() {
+ @Override
+ public void run(TransactionProxy transactionProxy) {
+ doReturn(incompleteFuture()).when(mockActorContext).executeOperationAsync(
+ any(ActorSelection.class), eqDataExists());
+
+ transactionProxy.exists(TestModel.TEST_PATH);
+
+ transactionProxy.exists(TestModel.TEST_PATH);
+ }
+ });
+ }
+
+ @Test
+ public void testExistsThrottlingWhenShardNotFound(){
+
+ completeOperation(new TransactionProxyOperation() {
+ @Override
+ public void run(TransactionProxy transactionProxy) {
+ doReturn(incompleteFuture()).when(mockActorContext).executeOperationAsync(
+ any(ActorSelection.class), eqDataExists());
+
+ transactionProxy.exists(TestModel.TEST_PATH);
+
+ transactionProxy.exists(TestModel.TEST_PATH);
+ }
+ }, false);
+ }
+
+
+ @Test
+ public void testExistsCompletion(){
+ completeOperation(new TransactionProxyOperation() {
+ @Override
+ public void run(TransactionProxy transactionProxy) {
+ doReturn(dataExistsReply(true)).when(mockActorContext).executeOperationAsync(
+ any(ActorSelection.class), eqDataExists());
+
+ transactionProxy.exists(TestModel.TEST_PATH);
+
+ transactionProxy.exists(TestModel.TEST_PATH);
+ }
+ });
+
+ }
+
+ @Test
+ public void testExistsCompletionForLocalShard(){
+ final NormalizedNode nodeToRead = ImmutableNodes.containerNode(TestModel.TEST_QNAME);
+ completeOperationLocal(new TransactionProxyOperation() {
+ @Override
+ public void run(TransactionProxy transactionProxy) {
+ transactionProxy.exists(TestModel.TEST_PATH);
+
+ transactionProxy.exists(TestModel.TEST_PATH);
+ }
+ }, createDataTree(nodeToRead));
+
+ }
+
+ @Test
+ public void testExistsCompletionForLocalShardWhenExceptionOccurs(){
+ completeOperationLocal(new TransactionProxyOperation() {
+ @Override
+ public void run(TransactionProxy transactionProxy) {
+ transactionProxy.exists(TestModel.TEST_PATH);
+
+ transactionProxy.exists(TestModel.TEST_PATH);
+ }
+ }, createDataTree());