Implement removing trees and write some test cases for DataTreeStore

implementations/onblocks/datanodestore/DataNodeStore.cpp

@@ -79,6 +79,17 @@ uint64_t DataNodeStore::numNodes() const {
   return _blockstore->numBlocks();
 }
 
+void DataNodeStore::removeSubtree(unique_ptr<DataNode> node) {
+   DataInnerNode *inner = dynamic_cast<DataInnerNode*>(node.get());
+   if (inner != nullptr) {
+     for (int i = 0; i < inner->numChildren(); ++i) {
+       auto child = load(inner->getChild(i)->key());
+       removeSubtree(std::move(child));
+     }
+   }
+   remove(std::move(node));
+ }
+
 }
 }
 }

implementations/onblocks/datanodestore/DataNodeStore.h

@@ -36,8 +36,10 @@ public:
 
   void remove(std::unique_ptr<DataNode> node);
 
+  void removeSubtree(std::unique_ptr<DataNode> node);
+
   uint64_t numNodes() const;
-  //TODO Test overwriteNodeWith(), createNodeAsCopyFrom()
+  //TODO Test overwriteNodeWith(), createNodeAsCopyFrom(), removeSubtree()
 
 private:
   std::unique_ptr<DataNode> load(std::unique_ptr<blockstore::Block> block);

implementations/onblocks/datatreestore/DataTree.cpp

@@ -116,6 +116,10 @@ void DataTree::flush() const {
   _rootNode->flush();
 }
 
+unique_ptr<DataNode> DataTree::releaseRootNode() {
+  return std::move(_rootNode);
+}
+
 
 }
 }

implementations/onblocks/datatreestore/DataTree.h

@@ -30,10 +30,14 @@ public:
   const blockstore::Key &key() const;
 
   void flush() const;
+
 private:
   datanodestore::DataNodeStore *_nodeStore;
   std::unique_ptr<datanodestore::DataNode> _rootNode;
 
+  std::unique_ptr<datanodestore::DataNode> releaseRootNode();
+  friend class DataTreeStore;
+
   std::unique_ptr<datanodestore::DataLeafNode> addDataLeafAt(datanodestore::DataInnerNode *insertPos);
   cpputils::optional_ownership_ptr<datanodestore::DataNode> createChainOfInnerNodes(unsigned int num, datanodestore::DataLeafNode *leaf);
   std::unique_ptr<datanodestore::DataLeafNode> addDataLeafToFullTree();

implementations/onblocks/datatreestore/DataTreeStore.cpp

@@ -21,7 +21,12 @@ DataTreeStore::~DataTreeStore() {
 }
 
 unique_ptr<DataTree> DataTreeStore::load(const blockstore::Key &key) {
-  return make_unique<DataTree>(_nodeStore.get(), _nodeStore->load(key));
+  auto node = _nodeStore->load(key);
+  if (node.get() == nullptr) {
+    return nullptr;
+  } else {
+    return make_unique<DataTree>(_nodeStore.get(), std::move(node));
+  }
 }
 
 unique_ptr<DataTree> DataTreeStore::createNewTree() {
@@ -29,6 +34,12 @@ unique_ptr<DataTree> DataTreeStore::createNewTree() {
   return make_unique<DataTree>(_nodeStore.get(), std::move(newleaf));
 }
 
+void DataTreeStore::remove(unique_ptr<DataTree> tree) {
+  auto root = tree->releaseRootNode();
+  tree.reset();
+  _nodeStore->removeSubtree(std::move(root));
+}
+
 }
 }
 }

implementations/onblocks/datatreestore/DataTreeStore.h

@@ -25,6 +25,8 @@ public:
 
   std::unique_ptr<DataTree> createNewTree();
 
+  void remove(std::unique_ptr<DataTree> tree);
+
 private:
   std::unique_ptr<datanodestore::DataNodeStore> _nodeStore;
 

test/implementations/onblocks/datanodestore/DataNodeStoreTest.cpp

@@ -96,9 +96,9 @@ TEST_F(DataNodeStoreTest, CreatedLeafNodeIsInitialized) {
 TEST_F(DataNodeStoreTest, NodeIsNotLoadableAfterDeleting) {
   auto nodekey = nodeStore->createNewLeafNode()->key();
   auto node = nodeStore->load(nodekey);
-  EXPECT_NE(nullptr, node);
+  EXPECT_NE(nullptr, node.get());
   nodeStore->remove(std::move(node));
-  EXPECT_EQ(nullptr, nodeStore->load(nodekey));
+  EXPECT_EQ(nullptr, nodeStore->load(nodekey).get());
 }
 
 TEST_F(DataNodeStoreTest, NumNodesIsCorrectOnEmptyNodestore) {

test/implementations/onblocks/datatreestore/DataTreeStoreTest.cpp

@@ -0,0 +1,44 @@
+#include "testutils/DataTreeTest.h"
+
+#include "../../../../implementations/onblocks/datanodestore/DataNodeStore.h"
+#include "../../../../implementations/onblocks/datatreestore/DataTreeStore.h"
+#include <messmer/blockstore/implementations/testfake/FakeBlockStore.h>
+
+using blockstore::testfake::FakeBlockStore;
+using blockstore::Key;
+using blobstore::onblocks::datanodestore::DataNodeStore;
+using std::make_unique;
+
+using namespace blobstore::onblocks::datatreestore;
+
+class DataTreeStoreTest: public DataTreeTest {
+};
+
+TEST_F(DataTreeStoreTest, NewTreeIsLeafOnly) {
+  auto tree = treeStore.createNewTree();
+  tree->flush();
+
+  EXPECT_IS_LEAF_NODE(tree->key());
+}
+
+TEST_F(DataTreeStoreTest, TreeIsNotLoadableAfterRemove) {
+  Key key = treeStore.createNewTree()->key();
+  auto tree = treeStore.load(key);
+  EXPECT_NE(nullptr, tree.get());
+  treeStore.remove(std::move(tree));
+  EXPECT_EQ(nullptr, treeStore.load(key).get());
+}
+
+TEST_F(DataTreeStoreTest, RemovingTreeRemovesAllNodesOfTheTree) {
+  auto key = CreateFullThreeLevel()->key();
+  auto tree1 = treeStore.load(key);
+  auto tree2_key = treeStore.createNewTree()->key();
+
+  treeStore.remove(std::move(tree1));
+
+  //Check that the only remaining node is tree2
+  EXPECT_EQ(1, nodeStore->numNodes());
+  EXPECT_NE(nullptr, treeStore.load(tree2_key).get());
+}
+
+//TODO ...

test/implementations/onblocks/datatreestore/growing/DataTreeGrowingTest_DataStaysIntact.cpp

@@ -18,7 +18,9 @@ class DataTreeGrowingTest_DataStaysIntact: public DataTreeGrowingTest {
 public:
   unique_ptr<DataTree> TreeWithData(unique_ptr<DataNode> root, TwoLevelDataFixture *data) {
     data->FillInto(root.get());
-    return make_unique<DataTree>(&nodeStore, std::move(root));
+    Key key = root->key();
+    root.reset();
+    return treeStore.load(key);
   }
 
   void TestDataStaysIntactOnGrowing(unique_ptr<DataNode> root, TwoLevelDataFixture *data) {
@@ -27,7 +29,7 @@ public:
     tree->addDataLeaf();
     tree->flush();
 
-    data->EXPECT_DATA_CORRECT(nodeStore.load(tree->key()).get(), numLeaves);
+    data->EXPECT_DATA_CORRECT(nodeStore->load(tree->key()).get(), numLeaves);
   }
 
   uint32_t countLeaves(DataNode *node) {
@@ -37,64 +39,64 @@ public:
     }
     uint32_t result = 0;
     for(int i = 0; i < inner->numChildren(); ++i) {
-      result += countLeaves(nodeStore.load(inner->getChild(i)->key()).get());
+      result += countLeaves(nodeStore->load(inner->getChild(i)->key()).get());
     }
     return result;
   }
 };
 
 TEST_F(DataTreeGrowingTest_DataStaysIntact, GrowAFullTwoLevelTree_FullLeaves) {
-  TwoLevelDataFixture data(&nodeStore, 0, true);
+  TwoLevelDataFixture data(nodeStore, 0, true);
   TestDataStaysIntactOnGrowing(CreateFullTwoLevel(), &data);
 }
 
 TEST_F(DataTreeGrowingTest_DataStaysIntact, GrowAFullTwoLevelTree_NonFullLeaves) {
-  TwoLevelDataFixture data(&nodeStore, 0, false);
+  TwoLevelDataFixture data(nodeStore, 0, false);
   TestDataStaysIntactOnGrowing(CreateFullTwoLevel(), &data);
 }
 
 TEST_F(DataTreeGrowingTest_DataStaysIntact, GrowAThreeLevelTreeWithLowerLevelFull_FullLeaves) {
-  TwoLevelDataFixture data(&nodeStore, 0, true);
+  TwoLevelDataFixture data(nodeStore, 0, true);
   auto node = CreateInner({CreateFullTwoLevel()});
   TestDataStaysIntactOnGrowing(std::move(node), &data);
 }
 
 TEST_F(DataTreeGrowingTest_DataStaysIntact, GrowAThreeLevelTreeWithLowerLevelFull_NonFullLeaves) {
-  TwoLevelDataFixture data(&nodeStore, 0, false);
+  TwoLevelDataFixture data(nodeStore, 0, false);
   auto node = CreateInner({CreateFullTwoLevel()});
   TestDataStaysIntactOnGrowing(std::move(node), &data);
 }
 
 TEST_F(DataTreeGrowingTest_DataStaysIntact, GrowAOneNodeTree_FullLeaves) {
-  TwoLevelDataFixture data(&nodeStore, 0, true);
+  TwoLevelDataFixture data(nodeStore, 0, true);
   TestDataStaysIntactOnGrowing(CreateLeaf(), &data);
 }
 
 TEST_F(DataTreeGrowingTest_DataStaysIntact, GrowAOneNodeTree_NonFullLeaves) {
-  TwoLevelDataFixture data(&nodeStore, 0, false);
+  TwoLevelDataFixture data(nodeStore, 0, false);
   TestDataStaysIntactOnGrowing(CreateLeaf(), &data);
 }
 
 TEST_F(DataTreeGrowingTest_DataStaysIntact, GrowATwoNodeTree_FullLeaves) {
-  TwoLevelDataFixture data(&nodeStore, 0, true);
+  TwoLevelDataFixture data(nodeStore, 0, true);
   auto node = CreateInner({CreateLeaf()});
   TestDataStaysIntactOnGrowing(std::move(node), &data);
 }
 
 TEST_F(DataTreeGrowingTest_DataStaysIntact, GrowATwoNodeTree_NonFullLeaves) {
-  TwoLevelDataFixture data(&nodeStore, 0, false);
+  TwoLevelDataFixture data(nodeStore, 0, false);
   auto node = CreateInner({CreateLeaf()});
   TestDataStaysIntactOnGrowing(std::move(node), &data);
 }
 
 TEST_F(DataTreeGrowingTest_DataStaysIntact, GrowAThreeNodeChainedTree_FullLeaves) {
-  TwoLevelDataFixture data(&nodeStore, 0, true);
+  TwoLevelDataFixture data(nodeStore, 0, true);
   auto node = CreateInner({CreateInner({CreateLeaf()})});
   TestDataStaysIntactOnGrowing(std::move(node), &data);
 }
 
 TEST_F(DataTreeGrowingTest_DataStaysIntact, GrowAThreeNodeChainedTree_NonFullLeaves) {
-  TwoLevelDataFixture data(&nodeStore, 0, false);
+  TwoLevelDataFixture data(nodeStore, 0, false);
   auto node = CreateInner({CreateInner({CreateLeaf()})});
   TestDataStaysIntactOnGrowing(std::move(node), &data);
 }

test/implementations/onblocks/datatreestore/growing/DataTreeGrowingTest_KeyDoesntChange.cpp

@@ -5,9 +5,9 @@ using blobstore::onblocks::datatreestore::DataTree;
 class DataTreeGrowingTest_KeyDoesntChange: public DataTreeGrowingTest {
 public:
   void EXPECT_KEY_DOESNT_CHANGE_WHEN_GROWING(const blockstore::Key &key) {
-    DataTree tree(&nodeStore, nodeStore.load(key));
-    tree.addDataLeaf();
-    EXPECT_EQ(key, tree.key());
+    auto tree = treeStore.load(key);
+    tree->addDataLeaf();
+    EXPECT_EQ(key, tree->key());
   }
 };
 

test/implementations/onblocks/datatreestore/growing/testutils/DataTreeGrowingTest.cpp

@@ -45,8 +45,7 @@ Key DataTreeGrowingTest::CreateThreeLevelTreeWithTwoFullSubtrees() {
 }
 
 void DataTreeGrowingTest::AddLeafTo(const Key &key) {
-  DataTree tree(&nodeStore, nodeStore.load(key));
-  tree.addDataLeaf();
+  treeStore.load(key)->addDataLeaf();
 }
 
 void DataTreeGrowingTest::EXPECT_IS_THREENODE_CHAIN(const Key &key) {

test/implementations/onblocks/datatreestore/impl/GetLowestInnerRightBorderNodeWithLessThanKChildrenOrNullTest.cpp

@@ -29,8 +29,8 @@ public:
   };
 
   void check(const TestData &testData) {
-    auto root = nodeStore.load(testData.rootNode);
-    auto result = GetLowestInnerRightBorderNodeWithLessThanKChildrenOrNull(&nodeStore, root.get());
+    auto root = nodeStore->load(testData.rootNode);
+    auto result = GetLowestInnerRightBorderNodeWithLessThanKChildrenOrNull(nodeStore, root.get());
     EXPECT_EQ(testData.expectedResult, result->key());
   }
 
@@ -58,8 +58,8 @@ public:
 };
 
 TEST_F(GetLowestInnerRightBorderNodeWithLessThanKChildrenOrNullTest, Leaf) {
-  auto leaf = nodeStore.createNewLeafNode();
-  auto result = GetLowestInnerRightBorderNodeWithLessThanKChildrenOrNull(&nodeStore, leaf.get());
+  auto leaf = nodeStore->createNewLeafNode();
+  auto result = GetLowestInnerRightBorderNodeWithLessThanKChildrenOrNull(nodeStore, leaf.get());
   EXPECT_EQ(nullptr, result.get());
 }
 
@@ -85,12 +85,12 @@ TEST_F(GetLowestInnerRightBorderNodeWithLessThanKChildrenOrNullTest, LargerTree)
 
 TEST_F(GetLowestInnerRightBorderNodeWithLessThanKChildrenOrNullTest, FullTwoLevelTree) {
   auto root = CreateFullTwoLevel();
-  auto result = GetLowestInnerRightBorderNodeWithLessThanKChildrenOrNull(&nodeStore, root.get());
+  auto result = GetLowestInnerRightBorderNodeWithLessThanKChildrenOrNull(nodeStore, root.get());
   EXPECT_EQ(nullptr, result.get());
 }
 
 TEST_F(GetLowestInnerRightBorderNodeWithLessThanKChildrenOrNullTest, FullThreeLevelTree) {
   auto root = CreateFullThreeLevel();
-  auto result = GetLowestInnerRightBorderNodeWithLessThanKChildrenOrNull(&nodeStore, root.get());
+  auto result = GetLowestInnerRightBorderNodeWithLessThanKChildrenOrNull(nodeStore, root.get());
   EXPECT_EQ(nullptr, result.get());
 }

test/implementations/onblocks/datatreestore/impl/GetLowestRightBorderNodeWithMoreThanOneChildOrNullTest.cpp

@@ -29,8 +29,8 @@ public:
   };
 
   void check(const TestData &testData) {
-    auto root = nodeStore.load(testData.rootNode);
-    auto result = GetLowestRightBorderNodeWithMoreThanOneChildOrNull(&nodeStore, root.get());
+    auto root = nodeStore->load(testData.rootNode);
+    auto result = GetLowestRightBorderNodeWithMoreThanOneChildOrNull(nodeStore, root.get());
     EXPECT_EQ(testData.expectedResult, result->key());
   }
 
@@ -81,20 +81,20 @@ public:
 };
 
 TEST_F(GetLowestRightBorderNodeWithMoreThanOneChildOrNullTest, Leaf) {
-  auto leaf = nodeStore.load(CreateLeafOnlyTree());
-  auto result = GetLowestRightBorderNodeWithMoreThanOneChildOrNull(&nodeStore, leaf.get());
+  auto leaf = nodeStore->load(CreateLeafOnlyTree());
+  auto result = GetLowestRightBorderNodeWithMoreThanOneChildOrNull(nodeStore, leaf.get());
   EXPECT_EQ(nullptr, result.get());
 }
 
 TEST_F(GetLowestRightBorderNodeWithMoreThanOneChildOrNullTest, TwoRightBorderNodes) {
-  auto node = nodeStore.load(CreateTwoRightBorderNodes());
-  auto result = GetLowestRightBorderNodeWithMoreThanOneChildOrNull(&nodeStore, node.get());
+  auto node = nodeStore->load(CreateTwoRightBorderNodes());
+  auto result = GetLowestRightBorderNodeWithMoreThanOneChildOrNull(nodeStore, node.get());
   EXPECT_EQ(nullptr, result.get());
 }
 
 TEST_F(GetLowestRightBorderNodeWithMoreThanOneChildOrNullTest, ThreeRightBorderNodes) {
-  auto node = nodeStore.load(CreateThreeRightBorderNodes());
-  auto result = GetLowestRightBorderNodeWithMoreThanOneChildOrNull(&nodeStore, node.get());
+  auto node = nodeStore->load(CreateThreeRightBorderNodes());
+  auto result = GetLowestRightBorderNodeWithMoreThanOneChildOrNull(nodeStore, node.get());
   EXPECT_EQ(nullptr, result.get());
 }
 

test/implementations/onblocks/datatreestore/shrinking/DataTreeShrinkingTest.cpp

@@ -17,7 +17,7 @@ using blockstore::Key;
 
 TEST_F(DataTreeShrinkingTest, ShrinkingALeafOnlyTreeCrashes) {
   Key key = CreateLeafOnlyTree()->key();
-  auto tree = make_unique<DataTree>(&nodeStore, nodeStore.load(key));
+  auto tree = treeStore.load(key);
   EXPECT_DEATH(tree->removeLastDataLeaf(), "");
 }
 
@@ -36,7 +36,7 @@ TEST_F(DataTreeShrinkingTest, ShrinkATwoLeafTree_LastLeafBlockIsDeleted) {
   auto lastChildKey = LoadInnerNode(tree->key())->getChild(1)->key();
 
   tree->removeLastDataLeaf();
-  EXPECT_EQ(nullptr, nodeStore.load(lastChildKey));
+  EXPECT_EQ(nullptr, nodeStore->load(lastChildKey));
 }
 
 TEST_F(DataTreeShrinkingTest, ShrinkATwoLeafTree_IntermediateBlocksAreDeleted) {
@@ -45,5 +45,5 @@ TEST_F(DataTreeShrinkingTest, ShrinkATwoLeafTree_IntermediateBlocksAreDeleted) {
   auto firstChildKey = LoadInnerNode(tree->key())->getChild(0)->key();
 
   tree->removeLastDataLeaf();
-  EXPECT_EQ(nullptr, nodeStore.load(firstChildKey));
+  EXPECT_EQ(nullptr, nodeStore->load(firstChildKey));
 }

test/implementations/onblocks/datatreestore/shrinking/DataTreeShrinkingTest_DataStaysIntact.cpp

@@ -14,7 +14,9 @@ class DataTreeShrinkingTest_DataStaysIntact: public DataTreeShrinkingTest {
 public:
   unique_ptr<DataTree> TreeWithData(unique_ptr<DataNode> root, TwoLevelDataFixture *data) {
     data->FillInto(root.get());
-    return make_unique<DataTree>(&nodeStore, std::move(root));
+    Key key = root->key();
+    root.reset();
+    return treeStore.load(key);
   }
 
   void TestDataStaysIntactOnShrinking(unique_ptr<DataInnerNode> root, TwoLevelDataFixture *data) {
@@ -22,106 +24,106 @@ public:
     tree->removeLastDataLeaf();
     tree->flush();
 
-    data->EXPECT_DATA_CORRECT(nodeStore.load(tree->key()).get());
+    data->EXPECT_DATA_CORRECT(nodeStore->load(tree->key()).get());
   }
 };
 
 TEST_F(DataTreeShrinkingTest_DataStaysIntact, ShrinkATwoLeafTree_FullLeaves) {
-  TwoLevelDataFixture data(&nodeStore, 0, true);
+  TwoLevelDataFixture data(nodeStore, 0, true);
   TestDataStaysIntactOnShrinking(CreateTwoLeaf(), &data);
 }
 
 TEST_F(DataTreeShrinkingTest_DataStaysIntact, ShrinkATwoLeafTree_NonFullLeaves) {
-  TwoLevelDataFixture data(&nodeStore, 0, false);
+  TwoLevelDataFixture data(nodeStore, 0, false);
   TestDataStaysIntactOnShrinking(CreateTwoLeaf(), &data);
 }
 
 TEST_F(DataTreeShrinkingTest_DataStaysIntact, ShrinkAFourNodeThreeLeafTree_FullLeaves) {
-  TwoLevelDataFixture data(&nodeStore, 0, true);
+  TwoLevelDataFixture data(nodeStore, 0, true);
   TestDataStaysIntactOnShrinking(CreateFourNodeThreeLeaf(), &data);
 }
 
 TEST_F(DataTreeShrinkingTest_DataStaysIntact, ShrinkAFourNodeThreeLeafTree_NonFullLeaves) {
-  TwoLevelDataFixture data(&nodeStore, 0, false);
+  TwoLevelDataFixture data(nodeStore, 0, false);
   TestDataStaysIntactOnShrinking(CreateFourNodeThreeLeaf(), &data);
 }
 
 TEST_F(DataTreeShrinkingTest_DataStaysIntact, ShrinkATwoInnerNodeOneTwoLeavesTree_FullLeaves) {
-  TwoLevelDataFixture data(&nodeStore, 0, true);
+  TwoLevelDataFixture data(nodeStore, 0, true);
   TestDataStaysIntactOnShrinking(CreateTwoInnerNodeOneTwoLeaves(), &data);
 }
 
 TEST_F(DataTreeShrinkingTest_DataStaysIntact, ShrinkATwoInnerNodeOneTwoLeavesTree_NonFullLeaves) {
-  TwoLevelDataFixture data(&nodeStore, 0, false);
+  TwoLevelDataFixture data(nodeStore, 0, false);
   TestDataStaysIntactOnShrinking(CreateTwoInnerNodeOneTwoLeaves(), &data);
 }
 
 TEST_F(DataTreeShrinkingTest_DataStaysIntact, ShrinkATwoInnerNodeTwoOneLeavesTree_FullLeaves) {
-  TwoLevelDataFixture data(&nodeStore, 0, true);
+  TwoLevelDataFixture data(nodeStore, 0, true);
   TestDataStaysIntactOnShrinking(CreateTwoInnerNodeTwoOneLeaves(), &data);
 }
 
 TEST_F(DataTreeShrinkingTest_DataStaysIntact, ShrinkATwoInnerNodeTwoOneLeavesTree_NonFullLeaves) {
-  TwoLevelDataFixture data(&nodeStore, 0, false);
+  TwoLevelDataFixture data(nodeStore, 0, false);
   TestDataStaysIntactOnShrinking(CreateTwoInnerNodeTwoOneLeaves(), &data);
 }
 
 TEST_F(DataTreeShrinkingTest_DataStaysIntact, ShrinkAThreeLevelMinDataTree_FullLeaves) {
-  TwoLevelDataFixture data(&nodeStore, 0, true);
+  TwoLevelDataFixture data(nodeStore, 0, true);
   TestDataStaysIntactOnShrinking(CreateThreeLevelMinData(), &data);
 }
 
 TEST_F(DataTreeShrinkingTest_DataStaysIntact, ShrinkAThreeLevelMinDataTree_NonFullLeaves) {
-  TwoLevelDataFixture data(&nodeStore, 0, false);
+  TwoLevelDataFixture data(nodeStore, 0, false);
   TestDataStaysIntactOnShrinking(CreateThreeLevelMinData(), &data);
 }
 
 TEST_F(DataTreeShrinkingTest_DataStaysIntact, ShrinkAFourLevelMinDataTree_FullLeaves) {
-  TwoLevelDataFixture data(&nodeStore, 0, true);
+  TwoLevelDataFixture data(nodeStore, 0, true);
   TestDataStaysIntactOnShrinking(CreateFourLevelMinData(), &data);
 }
 
 TEST_F(DataTreeShrinkingTest_DataStaysIntact, ShrinkAFourLevelMinDataTree_NonFullLeaves) {
-  TwoLevelDataFixture data(&nodeStore, 0, false);
+  TwoLevelDataFixture data(nodeStore, 0, false);
   TestDataStaysIntactOnShrinking(CreateFourLevelMinData(), &data);
 }
 
 TEST_F(DataTreeShrinkingTest_DataStaysIntact, ShrinkAFourLevelTreeWithTwoSiblingLeaves1_FullLeaves) {
-  TwoLevelDataFixture data(&nodeStore, 0, true);
+  TwoLevelDataFixture data(nodeStore, 0, true);
   TestDataStaysIntactOnShrinking(CreateFourLevelWithTwoSiblingLeaves1(), &data);
 }
 
 TEST_F(DataTreeShrinkingTest_DataStaysIntact, ShrinkAFourLevelTreeWithTwoSiblingLeaves1_NonFullLeaves) {
-  TwoLevelDataFixture data(&nodeStore, 0, false);
+  TwoLevelDataFixture data(nodeStore, 0, false);
   TestDataStaysIntactOnShrinking(CreateFourLevelWithTwoSiblingLeaves1(), &data);
 }
 
 TEST_F(DataTreeShrinkingTest_DataStaysIntact, ShrinkAFourLevelTreeWithTwoSiblingLeaves2_FullLeaves) {
-  TwoLevelDataFixture data(&nodeStore, 0, true);
+  TwoLevelDataFixture data(nodeStore, 0, true);
   TestDataStaysIntactOnShrinking(CreateFourLevelWithTwoSiblingLeaves2(), &data);
 }
 
 TEST_F(DataTreeShrinkingTest_DataStaysIntact, ShrinkAFourLevelTreeWithTwoSiblingLeaves2_NonFullLeaves) {
-  TwoLevelDataFixture data(&nodeStore, 0, false);
+  TwoLevelDataFixture data(nodeStore, 0, false);
   TestDataStaysIntactOnShrinking(CreateFourLevelWithTwoSiblingLeaves2(), &data);
 }
 
 TEST_F(DataTreeShrinkingTest_DataStaysIntact, ShrinkATreeWithFirstChildOfRootFullThreelevelAndSecondChildMindataThreelevel_FullLeaves) {
-  TwoLevelDataFixture data(&nodeStore, 0, true);
+  TwoLevelDataFixture data(nodeStore, 0, true);
   TestDataStaysIntactOnShrinking(CreateWithFirstChildOfRootFullThreelevelAndSecondChildMindataThreelevel(), &data);
 }
 
 TEST_F(DataTreeShrinkingTest_DataStaysIntact, ShrinkATreeWithFirstChildOfRootFullThreelevelAndSecondChildMindataThreelevel_NonFullLeaves) {
-  TwoLevelDataFixture data(&nodeStore, 0, false);
+  TwoLevelDataFixture data(nodeStore, 0, false);
   TestDataStaysIntactOnShrinking(CreateWithFirstChildOfRootFullThreelevelAndSecondChildMindataThreelevel(), &data);
 }
 
 TEST_F(DataTreeShrinkingTest_DataStaysIntact, ShrinkAThreeLevelTreeWithThreeChildrenOfRoot_FullLeaves) {
-  TwoLevelDataFixture data(&nodeStore, 0, true);
+  TwoLevelDataFixture data(nodeStore, 0, true);
   TestDataStaysIntactOnShrinking(CreateThreeLevelWithThreeChildrenOfRoot(), &data);
 }
 
 TEST_F(DataTreeShrinkingTest_DataStaysIntact, ShrinkAThreeLevelTreeWithThreeChildrenOfRoot_NonFullLeaves) {
-  TwoLevelDataFixture data(&nodeStore, 0, false);
+  TwoLevelDataFixture data(nodeStore, 0, false);
   TestDataStaysIntactOnShrinking(CreateThreeLevelWithThreeChildrenOfRoot(), &data);
 }

test/implementations/onblocks/datatreestore/shrinking/DataTreeShrinkingTest_KeyDoesntChange.cpp

@@ -6,9 +6,9 @@ using blockstore::Key;
 class DataTreeShrinkingTest_KeyDoesntChange: public DataTreeShrinkingTest {
 public:
   void EXPECT_KEY_DOESNT_CHANGE_WHEN_SHRINKING(const Key &key) {
-    DataTree tree(&nodeStore, nodeStore.load(key));
-    tree.removeLastDataLeaf();
-    EXPECT_EQ(key, tree.key());
+    auto tree = treeStore.load(key);
+    tree->removeLastDataLeaf();
+    EXPECT_EQ(key, tree->key());
   }
 };
 

test/implementations/onblocks/datatreestore/shrinking/DataTreeShrinkingTest_Structure.cpp

@@ -135,7 +135,7 @@ TEST_F(DataTreeShrinkingTest_Structure, ShrinkAFullTwoLevelTreeDownToOneLeaf) {
     Shrink(key);
   }
   EXPECT_IS_LEAF_NODE(key);
-  EXPECT_EQ(1, nodeStore.numNodes());
+  EXPECT_EQ(1, nodeStore->numNodes());
 }
 
 TEST_F(DataTreeShrinkingTest_Structure, ShrinkAFullThreeLevelTreeDownToOneLeaf) {
@@ -144,5 +144,5 @@ TEST_F(DataTreeShrinkingTest_Structure, ShrinkAFullThreeLevelTreeDownToOneLeaf)
     Shrink(key);
   }
   EXPECT_IS_LEAF_NODE(key);
-  EXPECT_EQ(1, nodeStore.numNodes());
+  EXPECT_EQ(1, nodeStore->numNodes());
 }

test/implementations/onblocks/datatreestore/shrinking/testutils/DataTreeShrinkingTest.cpp

@@ -9,8 +9,7 @@ using blockstore::Key;
 using blobstore::onblocks::datatreestore::DataTree;
 
 void DataTreeShrinkingTest::Shrink(const Key &key) {
-  DataTree tree(&nodeStore, nodeStore.load(key));
-  tree.removeLastDataLeaf();
+  treeStore.load(key)->removeLastDataLeaf();
 }
 
 unique_ptr<DataInnerNode> DataTreeShrinkingTest::CreateFourNodeThreeLeaf() {

test/implementations/onblocks/datatreestore/testutils/DataTreeTest.cpp

@@ -17,11 +17,13 @@ using std::vector;
 using cpputils::dynamic_pointer_move;
 
 DataTreeTest::DataTreeTest()
-  :nodeStore(make_unique<FakeBlockStore>()) {
+  :_nodeStore(make_unique<DataNodeStore>(make_unique<FakeBlockStore>())),
+   nodeStore(_nodeStore.get()),
+   treeStore(std::move(_nodeStore)) {
 }
 
 unique_ptr<DataLeafNode> DataTreeTest::CreateLeaf() {
-  return nodeStore.createNewLeafNode();
+  return nodeStore->createNewLeafNode();
 }
 
 unique_ptr<DataInnerNode> DataTreeTest::CreateInner(initializer_list<unique_ptr<DataNode>> children) {
@@ -36,7 +38,7 @@ unique_ptr<DataInnerNode> DataTreeTest::CreateInner(initializer_list<const DataN
 
 unique_ptr<DataInnerNode> DataTreeTest::CreateInner(vector<const DataNode*> children) {
   assert(children.size() >= 1);
-  auto node = nodeStore.createNewInnerNode(**children.begin());
+  auto node = nodeStore->createNewInnerNode(**children.begin());
   for(auto child = children.begin()+1; child != children.end(); ++child) {
     node->addChild(**child);
   }
@@ -44,7 +46,8 @@ unique_ptr<DataInnerNode> DataTreeTest::CreateInner(vector<const DataNode*> chil
 }
 
 unique_ptr<DataTree> DataTreeTest::CreateLeafOnlyTree() {
-  return make_unique<DataTree>(&nodeStore, CreateLeaf());
+  auto key = CreateLeaf()->key();
+  return treeStore.load(key);
 }
 
 void DataTreeTest::FillNode(DataInnerNode *node) {
@@ -72,14 +75,14 @@ unique_ptr<DataInnerNode> DataTreeTest::CreateFullThreeLevel() {
 }
 
 unique_ptr<DataInnerNode> DataTreeTest::LoadInnerNode(const Key &key) {
-  auto node = nodeStore.load(key);
+  auto node = nodeStore->load(key);
   auto casted = dynamic_pointer_move<DataInnerNode>(node);
   EXPECT_NE(nullptr, casted.get()) << "Is not an inner node";
   return casted;
 }
 
 unique_ptr<DataLeafNode> DataTreeTest::LoadLeafNode(const Key &key) {
-  auto node = nodeStore.load(key);
+  auto node = nodeStore->load(key);
   auto casted =  dynamic_pointer_move<DataLeafNode>(node);
   EXPECT_NE(nullptr, casted.get()) << "Is not a leaf node";
   return casted;
@@ -90,7 +93,8 @@ unique_ptr<DataInnerNode> DataTreeTest::CreateTwoLeaf() {
 }
 
 unique_ptr<DataTree> DataTreeTest::CreateTwoLeafTree() {
-  return make_unique<DataTree>(&nodeStore, CreateTwoLeaf());
+  auto key = CreateTwoLeaf()->key();
+  return treeStore.load(key);
 }
 
 void DataTreeTest::EXPECT_IS_LEAF_NODE(const Key &key) {

test/implementations/onblocks/datatreestore/testutils/DataTreeTest.h

@@ -8,6 +8,7 @@
 #include "../../../../../implementations/onblocks/datanodestore/DataInnerNode.h"
 #include "../../../../../implementations/onblocks/datanodestore/DataLeafNode.h"
 #include "../../../../../implementations/onblocks/datatreestore/DataTree.h"
+#include "../../../../../implementations/onblocks/datatreestore/DataTreeStore.h"
 
 class DataTreeTest: public ::testing::Test {
 public:
@@ -25,11 +26,14 @@ public:
   void FillNodeTwoLevel(blobstore::onblocks::datanodestore::DataInnerNode *node);
   std::unique_ptr<blobstore::onblocks::datanodestore::DataInnerNode> CreateFullTwoLevel();
   std::unique_ptr<blobstore::onblocks::datanodestore::DataInnerNode> CreateFullThreeLevel();
-  blobstore::onblocks::datanodestore::DataNodeStore nodeStore;
 
   std::unique_ptr<blobstore::onblocks::datanodestore::DataInnerNode> LoadInnerNode(const blockstore::Key &key);
   std::unique_ptr<blobstore::onblocks::datanodestore::DataLeafNode> LoadLeafNode(const blockstore::Key &key);
 
+  std::unique_ptr<blobstore::onblocks::datanodestore::DataNodeStore> _nodeStore;
+  blobstore::onblocks::datanodestore::DataNodeStore *nodeStore;
+  blobstore::onblocks::datatreestore::DataTreeStore treeStore;
+
   void EXPECT_IS_LEAF_NODE(const blockstore::Key &key);
   void EXPECT_IS_INNER_NODE(const blockstore::Key &key);
   void EXPECT_IS_TWONODE_CHAIN(const blockstore::Key &key);