Add: LeafStore.remove_leaf method

2 files changed, 19 insertions, 0 deletions

diff --git a/obnamlib/fmt_ga/leaf_store.py b/obnamlib/fmt_ga/leaf_store.py
index 5d488e96..bd356847 100644
--- a/obnamlib/fmt_ga/leaf_store.py
+++ b/obnamlib/fmt_ga/leaf_store.py
@@ -27,6 +27,9 @@ class LeafStoreInterface(object):  # pragma: no cover
     def get_leaf(self, leaf_id):
         raise NotImplementedError()
 
+    def remove_leaf(self, leaf_id):
+        raise NotImplementedError()
+
     def flush(self):
         raise NotImplementedError()
 
@@ -45,6 +48,10 @@ class InMemoryLeafStore(LeafStoreInterface):
     def get_leaf(self, leaf_id):
         return self._leaves.get(leaf_id, None)
 
+    def remove_leaf(self, leaf_id):
+        if leaf_id in self._leaves:
+            del self._leaves[leaf_id]
+
     def flush(self):
         pass
 
@@ -65,5 +72,11 @@ class LeafStore(LeafStoreInterface):  # pragma: no cover
         leaf.from_dict(self._blob_store.get_blob(leaf_id))
         return leaf
 
+    def remove_leaf(self, leaf_id):
+        # FIXME: This is a bit ugly, since we need to break the
+        # bag/blob store abstraction.
+        bag_id, _ = obnamlib.parse_object_id(leaf_id)
+        self._blob_store._bag_store.remove_bag(bag_id)
+
     def flush(self):
         self._blob_store.flush()
diff --git a/obnamlib/fmt_ga/leaf_store_tests.py b/obnamlib/fmt_ga/leaf_store_tests.py
index 9d1a8f71..92f6a01b 100644
--- a/obnamlib/fmt_ga/leaf_store_tests.py
+++ b/obnamlib/fmt_ga/leaf_store_tests.py
@@ -40,6 +40,12 @@ class LeafStoreTests(object):
     def test_returns_None_if_leaf_is_missing(self):
         self.assertEqual(self.ls.get_leaf(42), None)
 
+    def test_removes_leaf(self):
+        leaf = {'foo': 'bar'}
+        leaf_id = self.ls.put_leaf(leaf)
+        self.ls.remove_leaf(leaf_id)
+        self.assertEqual(self.ls.get_leaf(leaf_id), None)
+
     def test_has_flush(self):
         self.assertEqual(self.ls.flush(), None)