Add methods to set and get client keyids in ClientList.

2 files changed, 44 insertions, 0 deletions

diff --git a/obnamlib/clientlist.py b/obnamlib/clientlist.py
index 2ab09f91..577889dc 100644
--- a/obnamlib/clientlist.py
+++ b/obnamlib/clientlist.py
@@ -38,6 +38,7 @@ class ClientList(obnamlib.RepositoryTree):
     
     # subkey values
     CLIENT_NAME = 0
+    KEYID = 1
     SUBKEY_MAX = 255
 
     def __init__(self, fs, node_size, upload_queue_size, lru_size, hooks):
@@ -110,3 +111,22 @@ class ClientList(obnamlib.RepositoryTree):
             key = self.key(client_name, client_id, self.CLIENT_NAME)
             self.tree.remove(key)
 
+    def get_client_keyid(self, client_name):
+        if self.init_forest() and self.forest.trees:
+            t = self.forest.trees[-1]
+            client_id = self.find_client_id(t, client_name)
+            if client_id is not None:
+                key = self.key(client_name, client_id, self.KEYID)
+                for k, v in t.lookup_range(key, key):
+                    return v
+        return None
+        
+    def set_client_keyid(self, client_name, keyid):
+        self.start_changes()
+        client_id = self.find_client_id(self.tree, client_name)
+        key = self.key(client_name, client_id, self.KEYID)
+        if keyid is None:
+            self.tree.remove_range(key, key)
+        else:
+            self.tree.insert(key, keyid)
+
diff --git a/obnamlib/clientlist_tests.py b/obnamlib/clientlist_tests.py
index 9c19aeaa..a98bb088 100644
--- a/obnamlib/clientlist_tests.py
+++ b/obnamlib/clientlist_tests.py
@@ -64,6 +64,15 @@ class ClientListTests(unittest.TestCase):
         self.list.add_client('foo')
         self.list.remove_client('foo')
         self.assertEqual(self.list.get_client_id('foo'), None)
+        
+    def test_removed_client_has_no_keys(self):
+        self.list.add_client('foo')
+        client_id = self.list.get_client_id('foo')
+        self.list.remove_client('foo')
+        minkey = self.list.key('foo', client_id, 0)
+        maxkey = self.list.key('foo', client_id, self.list.SUBKEY_MAX)
+        pairs = list(self.list.tree.lookup_range(minkey, maxkey))
+        self.assertEqual(pairs, [])
 
     def test_twice_added_client_exists_only_once(self):
         self.list.add_client('foo')
@@ -80,3 +89,18 @@ class ClientListTests(unittest.TestCase):
         self.assertNotEqual(self.list.get_client_id('bar'),
                             self.list.get_client_id('foo'))
 
+    def test_client_has_no_public_key_initially(self):
+        self.list.add_client('foo')
+        self.assertEqual(self.list.get_client_keyid('foo'), None)
+
+    def test_sets_client_keyid(self):
+        self.list.add_client('foo')
+        self.list.set_client_keyid('foo', 'cafebeef')
+        self.assertEqual(self.list.get_client_keyid('foo'), 'cafebeef')
+
+    def test_remove_client_keyid(self):
+        self.list.add_client('foo')
+        self.list.set_client_keyid('foo', 'cafebeef')
+        self.list.set_client_keyid('foo', None)
+        self.assertEqual(self.list.get_client_keyid('foo'), None)
+