1 files changed, 20 insertions, 4 deletions
diff --git a/storage/src/tests/distributor/distributortestutil.cpp b/storage/src/tests/distributor/distributortestutil.cpp
index 2d2038b7cd0..b472ac1284e 100644
--- a/storage/src/tests/distributor/distributortestutil.cpp
+++ b/storage/src/tests/distributor/distributortestutil.cpp
@@ -43,14 +43,24 @@ DistributorTestUtil::setupDistributor(int redundancy,
                                       uint32_t earlyReturn,
                                       bool requirePrimaryToBeWritten)
 {
+    setup_distributor(redundancy, nodeCount, lib::ClusterStateBundle(lib::ClusterState(systemState)), earlyReturn, requirePrimaryToBeWritten);
+}
+
+void
+DistributorTestUtil::setup_distributor(int redundancy,
+                                       int node_count,
+                                       const lib::ClusterStateBundle& state,
+                                       uint32_t early_return,
+                                       bool require_primary_to_be_written)
+{
     lib::Distribution::DistributionConfigBuilder config(
-            lib::Distribution::getDefaultDistributionConfig(redundancy, nodeCount).get());
+            lib::Distribution::getDefaultDistributionConfig(redundancy, node_count).get());
     config.redundancy = redundancy;
-    config.initialRedundancy = earlyReturn;
-    config.ensurePrimaryPersisted = requirePrimaryToBeWritten;
+    config.initialRedundancy = early_return;
+    config.ensurePrimaryPersisted = require_primary_to_be_written;
     auto distribution = std::make_shared<lib::Distribution>(config);
     _node->getComponentRegister().setDistribution(distribution);
-    enableDistributorClusterState(systemState);
+    enable_distributor_cluster_state(state);
     // This is for all intents and purposes a hack to avoid having the
     // distributor treat setting the distribution explicitly as a signal that
     // it should send RequestBucketInfo to all configured nodes.
@@ -427,4 +437,10 @@ DistributorTestUtil::enableDistributorClusterState(vespalib::stringref state)
             lib::ClusterStateBundle(lib::ClusterState(state)));
 }
 
+void
+DistributorTestUtil::enable_distributor_cluster_state(const lib::ClusterStateBundle& state)
+{
+    getBucketDBUpdater().simulate_cluster_state_bundle_activation(state);
+}
+
 }