+ @Override
+ protected void pauseLeader(final Runnable operation) {
+ if (pauseLeaderFunction != null) {
+ pauseLeaderFunction.apply(operation);
+ } else {
+ super.pauseLeader(operation);
+ }
+ }
+
+ public static List<Object> fromState(final Snapshot.State from) {
+ if (from instanceof MockSnapshotState mockState) {
+ return mockState.getState();
+ }
+
+ throw new IllegalStateException("Unexpected snapshot State: " + from);
+ }
+
+ public ReplicatedLog getReplicatedLog() {
+ return getRaftActorContext().getReplicatedLog();
+ }
+
+ @Override
+ public Snapshot getRestoreFromSnapshot() {
+ return restoreFromSnapshot;
+ }
+
+ public static Props props(final String id, final Map<String, String> peerAddresses, final ConfigParams config) {
+ return builder().id(id).peerAddresses(peerAddresses).config(config).props();
+ }
+
+ public static Props props(final String id, final Map<String, String> peerAddresses,
+ final ConfigParams config, final DataPersistenceProvider dataPersistenceProvider) {
+ return builder().id(id).peerAddresses(peerAddresses).config(config)
+ .dataPersistenceProvider(dataPersistenceProvider).props();
+ }
+
+ public static Builder builder() {
+ return new Builder();
+ }
+
+ public static class AbstractBuilder<T extends AbstractBuilder<T, A>, A extends MockRaftActor> {
+ private Map<String, String> peerAddresses = Collections.emptyMap();
+ private String id;
+ private ConfigParams config;
+ private DataPersistenceProvider dataPersistenceProvider;
+ private ActorRef roleChangeNotifier;
+ private RaftActorSnapshotMessageSupport snapshotMessageSupport;
+ private Snapshot restoreFromSnapshot;
+ private Optional<Boolean> persistent = Optional.empty();
+ private final Class<A> actorClass;
+ private Function<Runnable, Void> pauseLeaderFunction;
+ private RaftActorSnapshotCohort snapshotCohort;
+
+ protected AbstractBuilder(final Class<A> actorClass) {
+ this.actorClass = actorClass;
+ }
+
+ @SuppressWarnings("unchecked")
+ private T self() {
+ return (T) this;
+ }
+
+ public T id(final String newId) {
+ id = newId;
+ return self();
+ }
+
+ public T peerAddresses(final Map<String, String> newPeerAddresses) {
+ peerAddresses = newPeerAddresses;
+ return self();
+ }
+
+ public T config(final ConfigParams newConfig) {
+ config = newConfig;
+ return self();
+ }
+
+ public T dataPersistenceProvider(final DataPersistenceProvider newDataPersistenceProvider) {
+ dataPersistenceProvider = newDataPersistenceProvider;
+ return self();
+ }
+
+ public T roleChangeNotifier(final ActorRef newRoleChangeNotifier) {
+ roleChangeNotifier = newRoleChangeNotifier;
+ return self();
+ }
+
+ public T snapshotMessageSupport(final RaftActorSnapshotMessageSupport newSnapshotMessageSupport) {
+ snapshotMessageSupport = newSnapshotMessageSupport;
+ return self();
+ }
+
+ public T restoreFromSnapshot(final Snapshot newRestoreFromSnapshot) {
+ restoreFromSnapshot = newRestoreFromSnapshot;
+ return self();
+ }
+
+ public T persistent(final Optional<Boolean> newPersistent) {
+ persistent = newPersistent;
+ return self();
+ }
+
+ public T pauseLeaderFunction(final Function<Runnable, Void> newPauseLeaderFunction) {
+ pauseLeaderFunction = newPauseLeaderFunction;
+ return self();
+ }
+
+ public T snapshotCohort(final RaftActorSnapshotCohort newSnapshotCohort) {
+ snapshotCohort = newSnapshotCohort;
+ return self();
+ }
+
+ public Props props() {
+ return Props.create(actorClass, this);
+ }
+ }
+
+ public static class Builder extends AbstractBuilder<Builder, MockRaftActor> {
+ Builder() {
+ super(MockRaftActor.class);
+ }
+ }
+
+ public static class MockSnapshotState implements Snapshot.State {
+ private static final long serialVersionUID = 1L;
+
+ private final List<Object> state;
+
+ public MockSnapshotState(final List<Object> state) {
+ this.state = state;
+ }
+
+ public List<Object> getState() {
+ return state;
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(state);
+ }
+
+ @Override
+ public boolean equals(final Object obj) {
+ if (this == obj) {
+ return true;
+ }
+ if (obj == null) {
+ return false;
+ }
+ if (getClass() != obj.getClass()) {
+ return false;
+ }
+ MockSnapshotState other = (MockSnapshotState) obj;
+ if (!Objects.equals(state, other.state)) {
+ return false;
+ }
+ return true;
+ }
+
+ @Override
+ public String toString() {
+ return "MockSnapshotState [state=" + state + "]";
+ }