Basic options for shell --build to use buildtreeswillsalmon/shellBuildTrees

Fixes issue #740

6 files changed, 175 insertions, 16 deletions

diff --git a/buildstream/_frontend/cli.py b/buildstream/_frontend/cli.py
index 68310cb15..f59038272 100644
--- a/buildstream/_frontend/cli.py
+++ b/buildstream/_frontend/cli.py
@@ -527,11 +527,14 @@ def show(app, elements, deps, except_, order, format_):
               help="Mount a file or directory into the sandbox")
 @click.option('--isolate', is_flag=True, default=False,
               help='Create an isolated build sandbox')
+@click.option('--use-buildtree', '-t', 'cli_buildtree', type=click.Choice(['ask', 'if-available', 'always', 'never']),
+              default='ask',
+              help='Defaults to ask but if set to always the function will fail if a build tree is not available')
 @click.argument('element', required=False,
                 type=click.Path(readable=False))
 @click.argument('command', type=click.STRING, nargs=-1)
 @click.pass_obj
-def shell(app, element, sysroot, mount, isolate, build_, command):
+def shell(app, element, sysroot, mount, isolate, build_, cli_buildtree, command):
     """Run a command in the target element's sandbox environment
 
     This will stage a temporary sysroot for running the target
@@ -557,6 +560,8 @@ def shell(app, element, sysroot, mount, isolate, build_, command):
     else:
         scope = Scope.RUN
 
+    use_buildtree = False
+
     with app.initialized():
         if not element:
             element = app.context.guess_element()
@@ -570,12 +575,30 @@ def shell(app, element, sysroot, mount, isolate, build_, command):
             HostMount(path, host_path)
             for host_path, path in mount
         ]
+
+        cached = element._cached_buildtree()
+        if cli_buildtree == "always":
+            if cached:
+                use_buildtree = True
+            else:
+                raise AppError("No buildtree is cached but the use buildtree option was specified")
+        elif cli_buildtree == "never":
+            pass
+        elif cli_buildtree == "if_available":
+            use_buildtree = cached
+        else:
+            if app.interactive and cached:
+                use_buildtree = bool(click.confirm('Do you want to use the cached buildtree?'))
+        if use_buildtree and not element._cached_success():
+            click.echo("Warning: using a buildtree from a failed build.")
+
         try:
             exitcode = app.stream.shell(element, scope, prompt,
                                         directory=sysroot,
                                         mounts=mounts,
                                         isolate=isolate,
-                                        command=command)
+                                        command=command,
+                                        usebuildtree=use_buildtree)
         except BstError as e:
             raise AppError("Error launching shell: {}".format(e), detail=e.detail) from e
 
diff --git a/buildstream/_stream.py b/buildstream/_stream.py
index 04eb409a5..85ead3330 100644
--- a/buildstream/_stream.py
+++ b/buildstream/_stream.py
@@ -124,6 +124,7 @@ class Stream():
     #    mounts (list of HostMount): Additional directories to mount into the sandbox
     #    isolate (bool): Whether to isolate the environment like we do in builds
     #    command (list): An argv to launch in the sandbox, or None
+    #    usebuildtree (bool): Wheather to use a buildtree as the source.
     #
     # Returns:
     #    (int): The exit code of the launched shell
@@ -132,7 +133,8 @@ class Stream():
               directory=None,
               mounts=None,
               isolate=False,
-              command=None):
+              command=None,
+              usebuildtree=False):
 
         # Assert we have everything we need built, unless the directory is specified
         # in which case we just blindly trust the directory, using the element
@@ -147,7 +149,8 @@ class Stream():
                 raise StreamError("Elements need to be built or downloaded before staging a shell environment",
                                   detail="\n".join(missing_deps))
 
-        return element._shell(scope, directory, mounts=mounts, isolate=isolate, prompt=prompt, command=command)
+        return element._shell(scope, directory, mounts=mounts, isolate=isolate, prompt=prompt, command=command,
+                              usebuildtree=usebuildtree)
 
     # build()
     #
diff --git a/buildstream/element.py b/buildstream/element.py
index 918323a2f..e44a7cc32 100644
--- a/buildstream/element.py
+++ b/buildstream/element.py
@@ -1338,11 +1338,12 @@ class Element(Plugin):
     # is used to stage things by the `bst checkout` codepath
     #
     @contextmanager
-    def _prepare_sandbox(self, scope, directory, shell=False, integrate=True):
+    def _prepare_sandbox(self, scope, directory, shell=False, integrate=True, usebuildtree=False):
         # bst shell and bst checkout require a local sandbox.
         bare_directory = True if directory else False
         with self.__sandbox(directory, config=self.__sandbox_config, allow_remote=False,
                             bare_directory=bare_directory) as sandbox:
+            sandbox._usebuildtree = usebuildtree
 
             # Configure always comes first, and we need it.
             self.__configure_sandbox(sandbox)
@@ -1386,7 +1387,7 @@ class Element(Plugin):
         # Stage all sources that need to be copied
         sandbox_vroot = sandbox.get_virtual_directory()
         host_vdirectory = sandbox_vroot.descend(directory.lstrip(os.sep).split(os.sep), create=True)
-        self._stage_sources_at(host_vdirectory, mount_workspaces=mount_workspaces)
+        self._stage_sources_at(host_vdirectory, mount_workspaces=mount_workspaces, usebuildtree=sandbox._usebuildtree)
 
     # _stage_sources_at():
     #
@@ -1395,10 +1396,10 @@ class Element(Plugin):
     # Args:
     #     vdirectory (:class:`.storage.Directory`): A virtual directory object to stage sources into.
     #     mount_workspaces (bool): mount workspaces if True, copy otherwise
+    #     usebuildtree (bool): use a the elements build tree as its source.
     #
-    def _stage_sources_at(self, vdirectory, mount_workspaces=True):
+    def _stage_sources_at(self, vdirectory, mount_workspaces=True, usebuildtree=False):
         with self.timed_activity("Staging sources", silent_nested=True):
-
             if not isinstance(vdirectory, Directory):
                 vdirectory = FileBasedDirectory(vdirectory)
             if not vdirectory.is_empty():
@@ -1420,7 +1421,7 @@ class Element(Plugin):
                                                  .format(workspace.get_absolute_path())):
                             workspace.stage(temp_staging_directory)
                 # Check if we have a cached buildtree to use
-                elif self._cached_buildtree():
+                elif usebuildtree:
                     artifact_base, _ = self.__extract()
                     import_dir = os.path.join(artifact_base, 'buildtree')
                 else:
@@ -1850,13 +1851,15 @@ class Element(Plugin):
     #    isolate (bool): Whether to isolate the environment like we do in builds
     #    prompt (str): A suitable prompt string for PS1
     #    command (list): An argv to launch in the sandbox
+    #    usebuildtree (bool): Use the buildtree as its source
     #
     # Returns: Exit code
     #
     # If directory is not specified, one will be staged using scope
-    def _shell(self, scope=None, directory=None, *, mounts=None, isolate=False, prompt=None, command=None):
+    def _shell(self, scope=None, directory=None, *, mounts=None, isolate=False, prompt=None, command=None,
+               usebuildtree=False):
 
-        with self._prepare_sandbox(scope, directory, shell=True) as sandbox:
+        with self._prepare_sandbox(scope, directory, shell=True, usebuildtree=usebuildtree) as sandbox:
             environment = self.get_environment()
             environment = copy.copy(environment)
             flags = SandboxFlags.INTERACTIVE | SandboxFlags.ROOT_READ_ONLY
diff --git a/buildstream/sandbox/sandbox.py b/buildstream/sandbox/sandbox.py
index b21b350ff..6081b0437 100644
--- a/buildstream/sandbox/sandbox.py
+++ b/buildstream/sandbox/sandbox.py
@@ -116,6 +116,7 @@ class Sandbox():
         self.__env = None
         self.__mount_sources = {}
         self.__allow_real_directory = kwargs['allow_real_directory']
+        self.usebuildtree = False
 
         # Plugin ID for logging
         plugin = kwargs.get('plugin', None)
@@ -146,6 +147,7 @@ class Sandbox():
 
         self._output_directory = None
         self._vdir = None
+        self._usebuildtree = False
 
         # This is set if anyone requests access to the underlying
         # directory via get_directory.
diff --git a/tests/integration/build-tree.py b/tests/integration/build-tree.py
index 5abadaab3..661658c38 100644
--- a/tests/integration/build-tree.py
+++ b/tests/integration/build-tree.py
@@ -19,7 +19,9 @@ DATA_DIR = os.path.join(
 @pytest.mark.datafiles(DATA_DIR)
 @pytest.mark.skipif(IS_LINUX and not HAVE_BWRAP, reason='Only available with bubblewrap on Linux')
 def test_buildtree_staged(cli_integration, tmpdir, datafiles):
-    # i.e. tests that cached build trees are staged by `bst shell --build`
+    # We can only test the non interacitve case
+    # The non interactive case defaults to not using buildtrees
+    # for `bst shell --build`
     project = os.path.join(datafiles.dirname, datafiles.basename)
     element_name = 'build-shell/buildtree.bst'
 
@@ -27,15 +29,67 @@ def test_buildtree_staged(cli_integration, tmpdir, datafiles):
     res.assert_success()
 
     res = cli_integration.run(project=project, args=[
-        'shell', '--build', element_name, '--', 'grep', '-q', 'Hi', 'test'
+        'shell', '--build', element_name, '--', 'cat', 'test'
+    ])
+    res.assert_shell_error()
+
+
+@pytest.mark.datafiles(DATA_DIR)
+@pytest.mark.skipif(IS_LINUX and not HAVE_BWRAP, reason='Only available with bubblewrap on Linux')
+def test_buildtree_staged_forced_true(cli_integration, tmpdir, datafiles):
+    # Test that if we ask for a build tree it is there.
+    project = os.path.join(datafiles.dirname, datafiles.basename)
+    element_name = 'build-shell/buildtree.bst'
+
+    res = cli_integration.run(project=project, args=['build', element_name])
+    res.assert_success()
+
+    res = cli_integration.run(project=project, args=[
+        'shell', '--build', '--use-buildtree', 'always', element_name, '--', 'cat', 'test'
+    ])
+    res.assert_success()
+    assert 'Hi' in res.output
+
+
+@pytest.mark.datafiles(DATA_DIR)
+@pytest.mark.skipif(IS_LINUX and not HAVE_BWRAP, reason='Only available with bubblewrap on Linux')
+def test_buildtree_staged_forced_true(cli_integration, tmpdir, datafiles):
+    # Test that a build tree can be correctly detected.
+    project = os.path.join(datafiles.dirname, datafiles.basename)
+    element_name = 'build-shell/buildtree.bst'
+
+    res = cli_integration.run(project=project, args=['build', element_name])
+    res.assert_success()
+
+    res = cli_integration.run(project=project, args=[
+        'shell', '--build', '--use-buildtree', 'if-available', element_name, '--', 'cat', 'test'
     ])
     res.assert_success()
+    assert 'Hi' in res.output
+
+
+@pytest.mark.datafiles(DATA_DIR)
+@pytest.mark.skipif(IS_LINUX and not HAVE_BWRAP, reason='Only available with bubblewrap on Linux')
+def test_buildtree_staged_forced_false(cli_integration, tmpdir, datafiles):
+    # Test that if we ask not to have a build tree it is not there
+    project = os.path.join(datafiles.dirname, datafiles.basename)
+    element_name = 'build-shell/buildtree.bst'
+
+    res = cli_integration.run(project=project, args=['build', element_name])
+    res.assert_success()
+
+    res = cli_integration.run(project=project, args=[
+        'shell', '--build', '--use-buildtree', 'never', element_name, '--', 'cat', 'test'
+    ])
+    res.assert_shell_error()
+
+    assert 'Hi' not in res.output
 
 
 @pytest.mark.datafiles(DATA_DIR)
 @pytest.mark.skipif(IS_LINUX and not HAVE_BWRAP, reason='Only available with bubblewrap on Linux')
 def test_buildtree_from_failure(cli_integration, tmpdir, datafiles):
-    # i.e. test that on a build failure, we can still shell into it
+    # Test that we can use a build tree after a failure
     project = os.path.join(datafiles.dirname, datafiles.basename)
     element_name = 'build-shell/buildtree-fail.bst'
 
@@ -44,9 +98,10 @@ def test_buildtree_from_failure(cli_integration, tmpdir, datafiles):
 
     # Assert that file has expected contents
     res = cli_integration.run(project=project, args=[
-        'shell', '--build', element_name, '--', 'cat', 'test'
+        'shell', '--build', element_name, '--use-buildtree', 'always', '--', 'cat', 'test'
     ])
     res.assert_success()
+    assert "Warning: using a buildtree from a failed build" in res.output
     assert 'Hi' in res.output
 
 
@@ -80,6 +135,65 @@ def test_buildtree_pulled(cli, tmpdir, datafiles):
 
         # Check it's using the cached build tree
         res = cli.run(project=project, args=[
-            'shell', '--build', element_name, '--', 'grep', '-q', 'Hi', 'test'
+            'shell', '--build', element_name, '--use-buildtree', 'always', '--', 'cat', 'test'
         ])
         res.assert_success()
+
+
+# This test checks for correct behaviour if a buildtree is not present.
+@pytest.mark.datafiles(DATA_DIR)
+@pytest.mark.skipif(IS_LINUX and not HAVE_BWRAP, reason='Only available with bubblewrap on Linux')
+def test_buildtree_options(cli, tmpdir, datafiles):
+    project = os.path.join(datafiles.dirname, datafiles.basename)
+    element_name = 'build-shell/buildtree.bst'
+
+    with create_artifact_share(os.path.join(str(tmpdir), 'artifactshare')) as share:
+        # Build the element to push it to cache
+        cli.configure({
+            'artifacts': {'url': share.repo, 'push': True}
+        })
+        result = cli.run(project=project, args=['build', element_name])
+        result.assert_success()
+        assert cli.get_element_state(project, element_name) == 'cached'
+
+        # Discard the cache
+        cli.configure({
+            'artifacts': {'url': share.repo, 'push': True},
+            'artifactdir': os.path.join(cli.directory, 'artifacts2')
+        })
+        assert cli.get_element_state(project, element_name) != 'cached'
+
+        # Pull from cache, but do not include buildtrees.
+        result = cli.run(project=project, args=['pull', '--deps', 'all', element_name])
+        result.assert_success()
+
+        # The above is the simplest way I know to create a local cache without any buildtrees.
+
+        # Check it's not using the cached build tree
+        res = cli.run(project=project, args=[
+            'shell', '--build', element_name, '--use-buildtree', 'never', '--', 'cat', 'test'
+        ])
+        res.assert_shell_error()
+        assert 'Hi' not in res.output
+
+        # Check it's not correctly handling the lack of buildtree
+        res = cli.run(project=project, args=[
+            'shell', '--build', element_name, '--use-buildtree', 'if-available', '--', 'cat', 'test'
+        ])
+        res.assert_shell_error()
+        assert 'Hi' not in res.output
+        
+        # Check it's not using the cached build tree, default is to ask, and fall back to not
+        # for non interactive behavior
+        res = cli.run(project=project, args=[
+            'shell', '--build', element_name, '--', 'cat', 'test'
+        ])
+        res.assert_shell_error()
+        assert 'Hi' not in res.output
+
+        # Check it's using the cached build tree
+        res = cli.run(project=project, args=[
+            'shell', '--build', element_name, '--use-buildtree', 'always', '--', 'cat', 'test'
+        ])
+        res.assert_main_error(ErrorDomain.PROG_NOT_FOUND, None)
+        assert 'Hi' not in res.output
diff --git a/tests/testutils/runcli.py b/tests/testutils/runcli.py
index ce5864bdf..f94cec8ae 100644
--- a/tests/testutils/runcli.py
+++ b/tests/testutils/runcli.py
@@ -153,6 +153,20 @@ class Result():
         assert self.task_error_domain == error_domain, fail_message
         assert self.task_error_reason == error_reason, fail_message
 
+    # assert_shell_error()
+    #
+    # Asserts that the buildstream created a shell and that the task in the
+    # shell failed.
+    #
+    # Args:
+    #    fail_message (str): An optional message to override the automatic
+    #                        assertion error messages
+    # Raises:
+    #    (AssertionError): If any of the assertions fail
+    #
+    def assert_shell_error(self, fail_message=''):
+        assert self.exit_code == 1, fail_message
+
     # get_tracked_elements()
     #
     # Produces a list of element names on which tracking occurred