_stream.py: Simplify code in shell()

Use single scheduler run for pulling dependencies and buildtree.
Deduplicate cache checks and corresponding warning/error messages
with and without pulling enabled.

2 files changed, 38 insertions, 102 deletions

diff --git a/src/buildstream/_stream.py b/src/buildstream/_stream.py
index 766c473ae..234211c49 100644
--- a/src/buildstream/_stream.py
+++ b/src/buildstream/_stream.py
@@ -198,108 +198,52 @@ class Stream:
         if unique_id and element is None:
             element = Plugin._lookup(unique_id)
         else:
-            # We may need to fetch dependency artifacts if we're pulling the artifact
-            selection = _PipelineSelection.ALL if pull_ else _PipelineSelection.NONE
+            selection = _PipelineSelection.BUILD if scope == _Scope.BUILD else _PipelineSelection.RUN
 
             elements = self.load_selection((element,), selection=selection, use_artifact_config=True)
 
-            # last one will be the element we want to stage, previous ones are
-            # elements to try and pull
-            element = elements[-1]
-            pull_dependencies = elements[:-1] if pull_ else None
-
-            artifact_is_cached = element._cached()
-            buildtree_is_cached = element._cached_buildtree()
-            buildtree_exists = element._buildtree_exists()
-            can_attempt_pull = self._context.pull_buildtrees and pull_
-
-            if usebuildtree:
-                if buildtree_is_cached:
-                    pass
-                # If element is already cached, we can check the proto to see if the buildtree existed
-                elif artifact_is_cached:
-                    if not buildtree_exists:
-                        if usebuildtree == "always":
-                            # Exit early if it won't be possible to even fetch a buildtree with always option
-                            raise StreamError("Artifact was created without buildtree, unable to launch shell with it")
-                        usebuildtree = None
-                        self._message(
-                            MessageType.WARN, "Artifact created without buildtree, shell will be loaded without it"
-                        )
-                    elif can_attempt_pull:
-                        self._message(
-                            MessageType.WARN,
-                            "buildtree is not cached locally but did exist, will attempt to pull from available remotes",
-                        )
-                    else:
-                        if usebuildtree == "always":
-                            # Exit early if it won't be possible to perform a fetch as pull semantics aren't present
-                            raise StreamError(
-                                "Artifact has a buildtree but it isn't cached. Can be retried with --pull and pull-buildtrees configured"
-                            )
-                        usebuildtree = None
-                        self._message(
-                            MessageType.WARN, "buildtree is not cached locally, shell will be loaded without it"
-                        )
-                # If element isn't cached at all, we can't check the proto to see if it existed so can't exit early
-                elif can_attempt_pull:
-                    if usebuildtree == "always":
-                        self._message(
-                            MessageType.WARN,
-                            "Element is not cached so buildtree status unknown, will attempt to pull from available remotes",
-                        )
-                else:
-                    if usebuildtree == "always":
-                        # Exit early as there is no buildtree locally & can_attempt_pull is False
-                        raise StreamError(
-                            "Artifact not cached locally. Can be retried with --pull and pull-buildtrees configured"
-                        )
-                    usebuildtree = None
-                    self._message(MessageType.WARN, "buildtree is not cached locally, shell will be loaded without it")
+            # Get element to stage from `targets` list.
+            # If scope is BUILD, it will not be in the `elements` list.
+            assert len(self.targets) == 1
+            element = self.targets[0]
+            element._set_required(scope)
 
-            # Raise warning if the element is cached in a failed state
-            if usebuildtree and element._cached_failure():
-                self._message(MessageType.WARN, "using a buildtree from a failed build.")
+            if pull_:
+                self._scheduler.clear_queues()
+                self._add_queue(PullQueue(self._scheduler))
+                plan = self._pipeline.add_elements([element], elements)
+                self._enqueue_plan(plan)
+                self._run()
 
         missing_deps = [dep for dep in self._pipeline.dependencies([element], scope) if not dep._cached()]
         if missing_deps:
-            if not pull_dependencies:
-                raise StreamError(
-                    "Elements need to be built or downloaded before staging a shell environment",
-                    detail="\n".join(list(map(lambda x: x._get_full_name(), missing_deps))),
-                )
-            self._message(MessageType.INFO, "Attempting to fetch missing or incomplete artifacts")
-            self._scheduler.clear_queues()
-            self._add_queue(PullQueue(self._scheduler))
-            plan = self._pipeline.add_elements([element], missing_deps)
-            self._enqueue_plan(plan)
-            self._run()
+            raise StreamError(
+                "Elements need to be built or downloaded before staging a shell environment",
+                detail="\n".join(list(map(lambda x: x._get_full_name(), missing_deps))),
+            )
 
         buildtree = False
         # Check if we require a pull queue attempt, with given artifact state and context
         if usebuildtree:
             if not element._cached_buildtree():
-                require_buildtree = self._buildtree_pull_required([element])
-                # Attempt a pull queue for the given element if remote and context allow it
-                if require_buildtree:
-                    self._message(MessageType.INFO, "Attempting to fetch missing artifact buildtree")
-                    self._scheduler.clear_queues()
-                    self._add_queue(PullQueue(self._scheduler))
-                    self._enqueue_plan(require_buildtree)
-                    self._run()
-                    # Now check if the buildtree was successfully fetched
-                    if element._cached_buildtree():
-                        buildtree = True
-
-                if not buildtree:
-                    message = "Buildtree is not cached locally or in available remotes"
-                    if usebuildtree == "always":
-                        raise StreamError(message)
-
-                    self._message(MessageType.INFO, message + ", shell will be loaded without it")
+                remotes_message = " or in available remotes" if pull_ else ""
+                if not element._cached():
+                    message = "Artifact not cached locally" + remotes_message
+                elif element._buildtree_exists():
+                    message = "Buildtree is not cached locally" + remotes_message
+                else:
+                    message = "Artifact was created without buildtree"
+                if usebuildtree == "always":
+                    raise StreamError(message)
+
+                self._message(MessageType.WARN, message + ", shell will be loaded without it")
             else:
                 buildtree = True
 
+                # Raise warning if the element is cached in a failed state
+                if element._cached_failure():
+                    self._message(MessageType.WARN, "using a buildtree from a failed build.")
+
         # Ensure we have our sources if we are launching a build shell
         if scope == _Scope.BUILD and not buildtree:
             self._fetch([element])
diff --git a/tests/integration/shellbuildtrees.py b/tests/integration/shellbuildtrees.py
index bb3a758e1..f4f13268a 100644
--- a/tests/integration/shellbuildtrees.py
+++ b/tests/integration/shellbuildtrees.py
@@ -70,14 +70,14 @@ def test_buildtree_staged_warn_empty_cached(cli_integration, tmpdir, datafiles):
         project=project, args=["shell", "--build", "--use-buildtree", "always", element_name, "--", "cat", "test"]
     )
     res.assert_main_error(ErrorDomain.APP, None)
-    assert "Artifact was created without buildtree, unable to launch shell with it" in res.stderr
+    assert "Error launching shell: Artifact was created without buildtree" in res.stderr
 
     # Now attempt the same with the try option, this should not attempt to find a buildtree
     # and just launch the shell, however the cat should still fail.
     res = cli_integration.run(
         project=project, args=["shell", "--build", "--use-buildtree", "try", element_name, "--", "cat", "test"]
     )
-    assert "Artifact created without buildtree, shell will be loaded without it" in res.stderr
+    assert "Artifact was created without buildtree, shell will be loaded without it" in res.stderr
     assert "Hi" not in res.output
 
 
@@ -153,7 +153,7 @@ def test_buildtree_from_failure_option_never(cli_integration, tmpdir, datafiles)
         project=project, args=["shell", "--build", element_name, "--use-buildtree", "always", "--", "cat", "test"]
     )
     res.assert_main_error(ErrorDomain.APP, None)
-    assert "Artifact was created without buildtree, unable to launch shell with it" in res.stderr
+    assert "Error launching shell: Artifact was created without buildtree" in res.stderr
 
 
 @pytest.mark.datafiles(DATA_DIR)
@@ -257,7 +257,7 @@ def test_buildtree_options(cli, tmpdir, datafiles):
         )
         assert "Hi" not in res.output
         assert "Attempting to fetch missing artifact buildtrees" not in res.stderr
-        assert "WARNING buildtree is not cached locally, shell will be loaded without it" in res.stderr
+        assert "WARNING Buildtree is not cached locally" in res.stderr
 
         # Check correctly handling the lack of buildtree, with 'try' attempting and succeeding
         # to pull the buildtree as the user context allow the pulling of buildtrees and it is
@@ -277,7 +277,6 @@ def test_buildtree_options(cli, tmpdir, datafiles):
                 "test",
             ],
         )
-        assert "Attempting to fetch missing artifact buildtree" in res.stderr
         assert "Hi" in res.output
         shutil.rmtree(os.path.join(os.path.join(str(tmpdir), "cache", "cas")))
         shutil.rmtree(os.path.join(os.path.join(str(tmpdir), "cache", "artifacts")))
@@ -291,12 +290,8 @@ def test_buildtree_options(cli, tmpdir, datafiles):
             project=project, args=["shell", "--build", element_name, "--use-buildtree", "always", "--", "cat", "test"]
         )
         res.assert_main_error(ErrorDomain.APP, None)
-        assert (
-            "Artifact has a buildtree but it isn't cached. Can be retried with --pull and pull-buildtrees configured"
-            in res.stderr
-        )
+        assert "Buildtree is not cached locally" in res.stderr
         assert "Hi" not in res.output
-        assert "Attempting to fetch missing artifact buildtree" not in res.stderr
 
         # Check that when user context is set to pull buildtrees and a remote has the buildtree,
         # 'always' will attempt and succeed at pulling the missing buildtree with --pull set.
@@ -316,10 +311,7 @@ def test_buildtree_options(cli, tmpdir, datafiles):
             ],
         )
         assert "Hi" in res.output
-        assert (
-            "buildtree is not cached locally but did exist, will attempt to pull from available remotes" in res.stderr
-        )
-        assert "Attempting to fetch missing artifact buildtree" in res.stderr
+        assert res.get_pulled_elements() == [element_name]
 
 
 # Tests running pull and pull-buildtree options at the same time.
@@ -348,7 +340,7 @@ def test_pull_buildtree_pulled(cli, tmpdir, datafiles):
             args=["shell", "--build", element_name, "--pull", "--use-buildtree", "always", "--", "cat", "test",],
         )
         res.assert_main_error(ErrorDomain.APP, None)
-        assert "Artifact not cached locally. Can be retried with --pull and pull-buildtrees configured" in res.stderr
+        assert "Buildtree is not cached locally" in res.stderr
 
         # Check it's using the cached build tree, because --pull
         # and pull-buildtrees were both set