Add: removal of keys from a CowTree

2 files changed, 36 insertions, 4 deletions

diff --git a/obnamlib/fmt_ga/cowtree.py b/obnamlib/fmt_ga/cowtree.py
index 69b00a9e..c4793a67 100644
--- a/obnamlib/fmt_ga/cowtree.py
+++ b/obnamlib/fmt_ga/cowtree.py
@@ -85,11 +85,22 @@ class CowTree(object):
         new_id = self._store.put_leaf(new)
         self._leaf_list.insert_leaf(sorted_keys[0], sorted_keys[-1], new_id)
 
-    def keys(self):
-        for leaf_id in self._leaf_list.leaf_ids():
+    def remove(self, key):
+        leaf_id = self._leaf_list.find_leaf_for_key(key)
+        if leaf_id is not None:
+            self._leaf_list.drop_leaf(leaf_id)
             leaf = self._store.get_leaf(leaf_id)
-            for key in leaf.keys():
-                yield key
+            leaf.remove(key)
+            if len(leaf) > 0:
+                self._make_split_leaf(leaf, list(sorted(leaf.keys())))
+
+    def keys(self):
+        leaf_ids = self._leaf_list.leaf_ids()
+        if leaf_ids:
+            for leaf_id in leaf_ids:
+                leaf = self._store.get_leaf(leaf_id)
+                for key in leaf.keys():
+                    yield key
 
     def commit(self):
         fake_leaf = obnamlib.CowLeaf()
diff --git a/obnamlib/fmt_ga/cowtree_tests.py b/obnamlib/fmt_ga/cowtree_tests.py
index c53a9ab7..2bcc4075 100644
--- a/obnamlib/fmt_ga/cowtree_tests.py
+++ b/obnamlib/fmt_ga/cowtree_tests.py
@@ -28,6 +28,9 @@ class CowTreeTests(unittest.TestCase):
         self.cow = obnamlib.CowTree()
         self.cow.set_leaf_store(self.ls)
 
+    def test_has_no_keys_initially(self):
+        self.assertEqual(list(self.cow.keys()), [])
+
     def test_lookup_returns_none_if_key_is_missing(self):
         self.assertEqual(self.cow.lookup(42), None)
 
@@ -36,6 +39,24 @@ class CowTreeTests(unittest.TestCase):
         value = 'barvalue'
         self.cow.insert(key, value)
         self.assertEqual(self.cow.lookup(key), value)
+        self.assertEqual(list(self.cow.keys()), [key])
+
+    def test_removes_only_key_in_blob(self):
+        key = 'fookey'
+        value = 'barvalue'
+        self.cow.insert(key, value)
+        self.cow.remove(key)
+        self.assertEqual(list(self.cow.keys()), [])
+
+    def test_removes_one_of_the_keys_in_blob(self):
+        key = 'fookey'
+        value = 'barvalue'
+        key2 = 'fookey2'
+        value2 = 'barvalue2'
+        self.cow.insert(key, value)
+        self.cow.insert(key2, value2)
+        self.cow.remove(key)
+        self.assertEqual(list(self.cow.keys()), [key2])
 
     def test_inserts_many_keys(self):
         N = 10