Make YAML file be able to specify what obnam commands to run

1 files changed, 24 insertions, 5 deletions

diff --git a/obbench b/obbench
index ffec6b3..8a84051 100755
--- a/obbench
+++ b/obbench
@@ -67,7 +67,8 @@ class ObnamBenchmarker(cliapp.Application):
         config = self.create_obnam_config(benchmark, tempdir)
         live = self.create_live_dir(tempdir)
         for step in benchmark.get('steps', []):
-            self.run_benchmark_step(step, checkout, config, live, result)
+            self.run_benchmark_step(
+                step, tempdir, checkout, config, live, result)
         return result
 
     def create_obnam_config(self, spec, tempdir):
@@ -88,18 +89,36 @@ class ObnamBenchmarker(cliapp.Application):
         os.mkdir(live)
         return live
 
-    def run_benchmark_step(self, step, checkout, config, live, result):
+    def run_benchmark_step(self,
+                           step, tempdir, checkout, config, live, result):
+        step_info = dict(step)
+
         if 'live' in step:
             cliapp.runcmd(['sh', '-euc', step['live']], cwd=live)
 
+        action = step['obnam']
+        func = funcs = {
+            'backup': self.run_backup,
+            'restore': self.run_restore,
+        }
         started = time.time()
-        self.run_obnam(checkout, ['backup', '--config', config])
+        funcs[action](tempdir, checkout, config)
         ended = time.time()
-
-        step_info = dict(step)
         step_info['duration'] = ended - started
+
         result.add_step(step_info)
 
+    def run_backup(self, tempdir, checkout, config):
+        self.run_obnam(checkout, ['backup', '--config', config])
+
+    def run_restore(self, tempdir, checkout, config):
+        restored = os.path.join(tempdir, 'restored')
+        if os.path.exists(restored):
+            shutil.rmtree(restored)
+        self.run_obnam(
+            checkout,
+            ['restore', '--config', config, '--to', restored])
+
     def run_obnam(self, checkout, args):
         cliapp.runcmd(
             ['./obnam', '--no-default-config'] + args,