1 files changed, 14 insertions, 2 deletions

diff --git a/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/ClusterFixture.java b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/ClusterFixture.java
index eccdb4d7832..f8d41405e85 100644
--- a/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/ClusterFixture.java
+++ b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/ClusterFixture.java
@@ -9,10 +9,13 @@ import com.yahoo.vdslib.state.NodeState;
 import com.yahoo.vdslib.state.NodeType;
 import com.yahoo.vdslib.state.State;
 import com.yahoo.vespa.clustercontroller.core.listeners.NodeListener;
+
 import java.util.Collection;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
+import java.util.TreeMap;
 
 import static org.mockito.Mockito.mock;
 
@@ -21,6 +24,7 @@ public class ClusterFixture {
     public final ContentCluster cluster;
     public final Distribution distribution;
     public final FakeTimer timer;
+    private final EventLogInterface eventLog;
     final StateChangeHandler nodeStateChangeHandler;
     private final ClusterStateGenerator.Params params = new ClusterStateGenerator.Params();
 
@@ -28,8 +32,9 @@ public class ClusterFixture {
         this.cluster = cluster;
         this.distribution = distribution;
         this.timer = new FakeTimer();
+        this.eventLog = mock(EventLogInterface.class);
         var context = new FleetControllerContextImpl(new FleetControllerId(cluster.getName(), 0));
-        this.nodeStateChangeHandler = new StateChangeHandler(context, timer, mock(EventLogInterface.class));
+        this.nodeStateChangeHandler = new StateChangeHandler(context, timer, eventLog);
         this.params.cluster(this.cluster);
     }
 
@@ -145,6 +150,13 @@ public class ClusterFixture {
         return this;
     }
 
+    static Map<NodeType, Integer> buildTransitionTimeMap(int distributorTransitionTime, int storageTransitionTime) {
+        Map<NodeType, Integer> maxTransitionTime = new TreeMap<>();
+        maxTransitionTime.put(NodeType.DISTRIBUTOR, distributorTransitionTime);
+        maxTransitionTime.put(NodeType.STORAGE, storageTransitionTime);
+        return maxTransitionTime;
+    }
+
     void disableTransientMaintenanceModeOnDown() {
         this.params.transitionTimes(0);
     }
@@ -162,7 +174,7 @@ public class ClusterFixture {
     }
 
     AnnotatedClusterState annotatedGeneratedClusterState() {
-        params.currentTimeInMillis(timer.getCurrentTimeInMillis());
+        params.currentTimeInMilllis(timer.getCurrentTimeInMillis());
         return ClusterStateGenerator.generatedStateFrom(params);
     }