Stop using old API to get distribution and bucket db.

2 files changed, 26 insertions, 14 deletions

diff --git a/storage/src/tests/distributor/bucketdbupdatertest.cpp b/storage/src/tests/distributor/bucketdbupdatertest.cpp
index 4a3f93d3f45..54e4c729c07 100644
--- a/storage/src/tests/distributor/bucketdbupdatertest.cpp
+++ b/storage/src/tests/distributor/bucketdbupdatertest.cpp
@@ -185,8 +185,9 @@ public:
             }
 
             std::vector<uint16_t> nodes;
-            getBucketDBUpdater().getDistributorComponent()
-                    .getDistribution().getIdealNodes(
+            auto &bucketSpaceRepo(getBucketDBUpdater().getDistributorComponent().getBucketSpaceRepo());
+            auto &distributorBucketSpace(bucketSpaceRepo.get(makeBucketSpace()));
+            distributorBucketSpace.getDistribution().getIdealNodes(
                     lib::NodeType::STORAGE,
                     state,
                     document::BucketId(16, i),
@@ -247,7 +248,9 @@ public:
         }
 
         std::vector<uint16_t> nodes;
-        getBucketDBUpdater().getDistributorComponent().getDistribution().getIdealNodes(
+        auto &bucketSpaceRepo(getBucketDBUpdater().getDistributorComponent().getBucketSpaceRepo());
+        auto &distributorBucketSpace(bucketSpaceRepo.get(makeBucketSpace()));
+        distributorBucketSpace.getDistribution().getIdealNodes(
                 lib::NodeType::STORAGE,
                 state,
                 document::BucketId(id),
@@ -584,8 +587,10 @@ BucketDBUpdaterTest::testNormalUsage()
     CPPUNIT_ASSERT_EQUAL(size_t(3), _sender.commands.size());
 
     // Ensure distribution hash is set correctly
+    auto &bucketSpaceRepo(getBucketDBUpdater().getDistributorComponent().getBucketSpaceRepo());
+    auto &distributorBucketSpace(bucketSpaceRepo.get(makeBucketSpace()));
     CPPUNIT_ASSERT_EQUAL(
-            getBucketDBUpdater().getDistributorComponent().getDistribution()
+            distributorBucketSpace.getDistribution()
             .getNodeGraph().getDistributionConfigHash(),
             dynamic_cast<const RequestBucketInfoCommand&>(
                     *_sender.commands[0]).getDistributionHash());
@@ -916,8 +921,9 @@ BucketDBUpdaterTest::testBitChange()
 
         int cnt=0;
         for (int i=0; cnt < 2; i++) {
-            lib::Distribution distribution = getBucketDBUpdater().getDistributorComponent()
-                            .getDistribution();
+            auto &bucketSpaceRepo(getBucketDBUpdater().getDistributorComponent().getBucketSpaceRepo());
+            auto &distributorBucketSpace(bucketSpaceRepo.get(makeBucketSpace()));
+            lib::Distribution distribution = distributorBucketSpace.getDistribution();
             std::vector<uint16_t> distributors;
             if (distribution.getIdealDistributorNode(
                     lib::ClusterState("redundancy:1 bits:14 storage:1 distributor:2"),
@@ -1835,10 +1841,11 @@ BucketDBUpdaterTest::mergeBucketLists(
     }
 
     BucketDumper dumper(includeBucketInfo);
-    getBucketDBUpdater().getDistributorComponent()
-            .getBucketDatabase().forEach(dumper);
-    getBucketDBUpdater().getDistributorComponent()
-            .getBucketDatabase().clear();
+    auto &bucketSpaceRepo(getBucketDBUpdater().getDistributorComponent().getBucketSpaceRepo());
+    auto &distributorBucketSpace(bucketSpaceRepo.get(makeBucketSpace()));
+    auto &bucketDb(distributorBucketSpace.getBucketDatabase());
+    bucketDb.forEach(dumper);
+    bucketDb.clear();
     return dumper.ost.str();
 }
 
diff --git a/storage/src/tests/distributor/operationtargetresolvertest.cpp b/storage/src/tests/distributor/operationtargetresolvertest.cpp
index 30891d17b3f..72f028a3ff8 100644
--- a/storage/src/tests/distributor/operationtargetresolvertest.cpp
+++ b/storage/src/tests/distributor/operationtargetresolvertest.cpp
@@ -10,6 +10,7 @@
 #include <tests/distributor/distributortestutil.h>
 #include <vespa/vdslib/distribution/idealnodecalculatorimpl.h>
 #include <vespa/vespalib/testkit/testapp.h>
+#include <vespa/storage/distributor/distributor_bucket_space_repo.h>
 #include <vespa/storage/distributor/operationtargetresolverimpl.h>
 #include <vespa/storage/distributor/externaloperationhandler.h>
 #include <vespa/config/helper/configgetter.hpp>
@@ -148,11 +149,13 @@ OperationTargetResolverTest::getInstances(const BucketId& id,
                                           bool stripToRedundancy)
 {
     lib::IdealNodeCalculatorImpl idealNodeCalc;
-    idealNodeCalc.setDistribution(getExternalOperationHandler().getDistribution());
+    auto &bucketSpaceRepo(getExternalOperationHandler().getBucketSpaceRepo());
+    auto &distributorBucketSpace(bucketSpaceRepo.get(makeBucketSpace()));
+    idealNodeCalc.setDistribution(distributorBucketSpace.getDistribution());
     idealNodeCalc.setClusterState(getExternalOperationHandler().getClusterState());
     OperationTargetResolverImpl resolver(
-            getExternalOperationHandler().getBucketDatabase(), idealNodeCalc, 16,
-            getExternalOperationHandler().getDistribution().getRedundancy(),
+            distributorBucketSpace.getBucketDatabase(), idealNodeCalc, 16,
+            distributorBucketSpace.getDistribution().getRedundancy(),
             makeBucketSpace());
     if (stripToRedundancy) {
         return resolver.getInstances(OperationTargetResolver::PUT, id);
@@ -179,11 +182,13 @@ OperationTargetResolverTest::testMultipleNodes()
 {
     setupDistributor(1, 2, "storage:2 distributor:1");
 
+    auto &bucketSpaceRepo(getExternalOperationHandler().getBucketSpaceRepo());
+    auto &distributorBucketSpace(bucketSpaceRepo.get(makeBucketSpace()));
     for (int i = 0; i < 100; ++i) {
         addNodesToBucketDB(BucketId(16, i), "0=0,1=0");
 
         lib::IdealNodeCalculatorImpl idealNodeCalc;
-        idealNodeCalc.setDistribution(getExternalOperationHandler().getDistribution());
+        idealNodeCalc.setDistribution(distributorBucketSpace.getDistribution());
         idealNodeCalc.setClusterState(getExternalOperationHandler().getClusterState());
         lib::IdealNodeList idealNodes(
                 idealNodeCalc.getIdealStorageNodes(BucketId(16, i)));