_frontend/cli.py: Add option to reset multiple workspaces

!357 added support for closing multiple workspaces. Similarly, also
allow `bst workspace reset` to work on multiple workspaces, with `--all`
as a helper to reset all open workspaces.

2 files changed, 68 insertions, 6 deletions

diff --git a/buildstream/_frontend/cli.py b/buildstream/_frontend/cli.py
index 628a6cc85..3d50d5c14 100644
--- a/buildstream/_frontend/cli.py
+++ b/buildstream/_frontend/cli.py
@@ -660,20 +660,30 @@ def workspace_close(app, remove_dir, all_, elements):
 @workspace.command(name='reset', short_help="Reset a workspace to its original state")
 @click.option('--track', 'track_', default=False, is_flag=True,
               help="Track and fetch the latest source before resetting")
-@click.argument('element',
+@click.option('--all', '-a', 'all_', default=False, is_flag=True,
+              help="Reset all open workspaces")
+@click.argument('elements', nargs=-1,
                 type=click.Path(dir_okay=False, readable=True))
 @click.pass_obj
-def workspace_reset(app, track_, element):
+def workspace_reset(app, track_, all_, elements):
     """Reset a workspace to its original state"""
+
+    if not (all_ or elements):
+        click.echo('ERROR: no elements specified', err=True)
+        sys.exit(-1)
+
     if app.interactive:
         if not click.confirm('This will remove all your changes, are you sure?'):
             click.echo('Aborting', err=True)
             sys.exit(-1)
 
-    with app.initialized((element,)):
-        # This command supports only one target
-        target = app.pipeline.targets[0]
-        app.reset_workspace(target, track_)
+    with app.partially_initialized():
+        if all_:
+            elements = tuple(element_name for element_name, _ in app.project.workspaces.list())
+
+    with app.initialized(elements):
+        for target in app.pipeline.targets:
+            app.reset_workspace(target, track_)
 
 
 ##################################################################
diff --git a/tests/frontend/workspace.py b/tests/frontend/workspace.py
index 6ecc9e2d2..e45696d38 100644
--- a/tests/frontend/workspace.py
+++ b/tests/frontend/workspace.py
@@ -237,6 +237,58 @@ def test_reset(cli, tmpdir, datafiles):
 
 
 @pytest.mark.datafiles(DATA_DIR)
+def test_reset_multiple(cli, tmpdir, datafiles):
+    # Open the workspaces
+    tmpdir_alpha = os.path.join(str(tmpdir), 'alpha')
+    tmpdir_beta = os.path.join(str(tmpdir), 'beta')
+    alpha, project, workspace_alpha = open_workspace(
+        cli, tmpdir_alpha, datafiles, 'git', False, suffix='-alpha')
+    beta, project, workspace_beta = open_workspace(
+        cli, tmpdir_beta, datafiles, 'git', False, suffix='-beta')
+
+    # Modify workspaces
+    shutil.rmtree(os.path.join(workspace_alpha, 'usr', 'bin'))
+    os.makedirs(os.path.join(workspace_beta, 'etc'))
+    with open(os.path.join(workspace_beta, 'etc', 'pony.conf'), 'w') as f:
+        f.write("PONY='pink'")
+
+    # Now reset the open workspaces, this should have the
+    # effect of reverting our changes.
+    result = cli.run(project=project, args=[
+        'workspace', 'reset', alpha, beta,
+    ])
+    result.assert_success()
+    assert os.path.exists(os.path.join(workspace_alpha, 'usr', 'bin', 'hello'))
+    assert not os.path.exists(os.path.join(workspace_beta, 'etc', 'pony.conf'))
+
+
+@pytest.mark.datafiles(DATA_DIR)
+def test_reset_all(cli, tmpdir, datafiles):
+    # Open the workspaces
+    tmpdir_alpha = os.path.join(str(tmpdir), 'alpha')
+    tmpdir_beta = os.path.join(str(tmpdir), 'beta')
+    alpha, project, workspace_alpha = open_workspace(
+        cli, tmpdir_alpha, datafiles, 'git', False, suffix='-alpha')
+    beta, project, workspace_beta = open_workspace(
+        cli, tmpdir_beta, datafiles, 'git', False, suffix='-beta')
+
+    # Modify workspaces
+    shutil.rmtree(os.path.join(workspace_alpha, 'usr', 'bin'))
+    os.makedirs(os.path.join(workspace_beta, 'etc'))
+    with open(os.path.join(workspace_beta, 'etc', 'pony.conf'), 'w') as f:
+        f.write("PONY='pink'")
+
+    # Now reset the open workspace, this should have the
+    # effect of reverting our changes.
+    result = cli.run(project=project, args=[
+        'workspace', 'reset', '--all'
+    ])
+    result.assert_success()
+    assert os.path.exists(os.path.join(workspace_alpha, 'usr', 'bin', 'hello'))
+    assert not os.path.exists(os.path.join(workspace_beta, 'etc', 'pony.conf'))
+
+
+@pytest.mark.datafiles(DATA_DIR)
 def test_list(cli, tmpdir, datafiles):
     element_name, project, workspace = open_workspace(cli, tmpdir, datafiles, 'git', False)