Add update_tree

1 files changed, 17 insertions, 0 deletions

diff --git a/israwhidebroken/client.py b/israwhidebroken/client.py
index 53c41ef..41f8dd1 100644
--- a/israwhidebroken/client.py
+++ b/israwhidebroken/client.py
@@ -62,6 +62,23 @@ class IRBClient(BaseClient):
         r = self.send_request('/add_tree', auth=True, req_params=kw)
         return r.get('tree')
 
+    def update_tree(self, *args, **kw):
+        '''Add a new tree, or update an existing tree.
+        Required arguments: id, or arch and compose_id
+        Optional: tree_time, repodata_time
+        Returns the tree data.'''
+        if 'id' in kw:
+            treelist = self.get_trees(id=id)
+        else:
+            treelist = self.get_trees(arch=kw['arch'],
+                                      compose_id=kw['compose_id'])
+        if len(treelist) == 0:
+            return self.add_tree(*args, **kw)
+
+        kw['treeid'] = treelist[0]['id']
+        r = self.send_request('/update_tree', auth=True, req_params=kw)
+        return r.get('tree')
+
     def add_result(self, treeid, testid, result, detail_url=None):
         '''Add a test result to the database. Returns result id on success.'''
         params = {'treeid':treeid, 'testid':testid, 'result': result}