- doThrow(exToThrow).when(mockActorContext).executeShardOperation(
- anyString(), any(), anyDuration());
+ if (exToThrow instanceof PrimaryNotFoundException) {
+ doReturn(Optional.absent()).when(mockActorContext).findPrimaryShard(anyString());
+ } else {
+ doReturn(Optional.of(getSystem().actorSelection(actorRef.path()))).
+ when(mockActorContext).findPrimaryShard(anyString());
+ }
+ doThrow(exToThrow).when(mockActorContext).executeOperation(any(ActorSelection.class), any());