build-providers: add support for --shell-after (#2253)

Adds support for shelling into the environment after the requested
step has been executed.

LP: #1782767

Signed-off-by: Sergio Schvezov <sergio.schvezov@canonical.com>

snapcraft/cli/_options.py

@@ -31,16 +31,18 @@ def get_help_record(self, ctx):
     "--no-parallel-builds",
     "--target-arch",
     "--debug",
+    "--shell-after",
 ]
 
 _BUILD_OPTIONS = [
     dict(
         is_flag=True,
-        help=("Detect best candidate location for stage-packages using geoip"),
+        help="Detect best candidate location for stage-packages using geoip",
     ),
     dict(is_flag=True, help="Force a sequential build."),
     dict(metavar="<arch>", help="Target architecture to cross compile to"),
     dict(is_flag=True, help="Shells into the environment if the build fails."),
+    dict(is_flag=True, help="Shells into the environment after the step has run."),
 ]
 
 

snapcraft/cli/lifecycle.py

@@ -45,6 +45,7 @@ def _execute(
     parts: str,
     pack_project: bool = False,
     output: str = None,
+    shell_after: bool = False,
     **kwargs
 ) -> "Project":
     # fmt: on
@@ -72,6 +73,9 @@ def _execute(
                     echo.warning("Run the same command again with --debug to shell into the environment "
                                  "if you wish to introspect this failure.")
                     raise
+            else:
+                if shell_after:
+                    instance.shell()
     elif build_environment.is_managed_host or build_environment.is_host:
         project_config = project_loader.load_config(project)
         lifecycle.execute(step, project_config, parts)
@@ -207,6 +211,7 @@ def pack(directory, output, **kwargs):
 
 
 @lifecyclecli.command()
+@add_build_options()
 @click.argument("parts", nargs=-1, metavar="<part>...", required=False)
 @click.option(
     "--step",
@@ -215,7 +220,7 @@ def pack(directory, output, **kwargs):
     type=click.Choice(["pull", "build", "stage", "prime", "strip"]),
     help="only clean the specified step and those that depend on it.",
 )
-def clean(parts, step_name):
+def clean(parts, step_name, **kwargs):
     """Remove content - cleans downloads, builds or install artifacts.
 
     \b
@@ -226,13 +231,14 @@ def clean(parts, step_name):
     build_environment = env.BuilderEnvironmentConfig()
     try:
         project = get_project(
-            is_managed_host=build_environment.is_managed_host
+            is_managed_host=build_environment.is_managed_host, **kwargs
         )
     except YamlValidationError:
         # We need to be able to clean invalid projects too.
         project = get_project(
             is_managed_host=build_environment.is_managed_host,
-            skip_snapcraft_yaml=True
+            skip_snapcraft_yaml=True,
+            **kwargs,
         )
 
     step = None
@@ -269,7 +275,10 @@ def clean(parts, step_name):
     metavar="<remote>",
     help="Use a specific lxd remote instead of a local container.",
 )
-def cleanbuild(remote, **kwargs):
+@click.option(
+    "--debug", is_flag=True, help="Shells into the environment if the build fails."
+)
+def cleanbuild(remote, debug, **kwargs):
     """Create a snap using a clean environment managed by a build provider.
 
     \b
@@ -297,7 +306,7 @@ def cleanbuild(remote, **kwargs):
         default=default_provider, additional_providers=["multipass"]
     )
     project = get_project(
-        is_managed=build_environment.is_managed_host, **kwargs
+        is_managed=build_environment.is_managed_host, **kwargs, debug=debug
     )
 
     snap_filename = lifecycle.cleanbuild(

tests/unit/commands/test_build_providers.py

@@ -15,6 +15,7 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 from unittest import mock
+from typing import Optional
 
 import fixtures
 from testtools.matchers import Equals
@@ -69,3 +70,72 @@ def test_step_without_debug_using_build_provider_fails_and_does_not_shell(self):
         self.assertRaises(ProviderExecError, self.run_command, ["pull"])
 
         self.shell_mock.assert_not_called()
+
+
+class BuildProviderShellCommandTestCase(LifecycleCommandsBaseTestCase):
+    def setUp(self):
+        super().setUp()
+        self.useFixture(
+            fixtures.EnvironmentVariable("SNAPCRAFT_BUILD_ENVIRONMENT", "multipass")
+        )
+
+        shell_mock = mock.Mock()
+        pack_project_mock = mock.Mock()
+        execute_step_mock = mock.Mock()
+
+        class Provider(ProviderImpl):
+            def pack_project(self, *, output: Optional[str] = None) -> None:
+                pack_project_mock(output)
+
+            def execute_step(self, step: steps.Step) -> None:
+                execute_step_mock(step)
+
+            def shell(self):
+                shell_mock()
+
+        patcher = mock.patch(
+            "snapcraft.internal.build_providers.get_provider_for", return_value=Provider
+        )
+        self.provider = patcher.start()
+        self.addCleanup(patcher.stop)
+
+        self.shell_mock = shell_mock
+        self.pack_project_mock = pack_project_mock
+        self.execute_step_mock = execute_step_mock
+
+        self.make_snapcraft_yaml("pull", base="core18")
+
+    def test_step_with_shell_after(self):
+        result = self.run_command(["pull", "--shell-after"])
+
+        self.assertThat(result.exit_code, Equals(0))
+        self.pack_project_mock.assert_not_called()
+        self.execute_step_mock.assert_called_once_with(steps.PULL)
+        self.shell_mock.assert_called_once_with()
+
+    def test_snap_with_shell_after(self):
+        result = self.run_command(["snap", "--output", "fake.snap", "--shell-after"])
+
+        self.assertThat(result.exit_code, Equals(0))
+        self.pack_project_mock.assert_called_once_with("fake.snap")
+        self.execute_step_mock.assert_not_called()
+        self.shell_mock.assert_called_once_with()
+
+    def test_step_without_shell_after_does_not_enter_shell(self):
+        result = self.run_command(["pull"])
+
+        self.assertThat(result.exit_code, Equals(0))
+        self.pack_project_mock.assert_not_called()
+        self.execute_step_mock.assert_called_once_with(steps.PULL)
+        self.shell_mock.assert_not_called()
+
+    def test_error_with_shell_after_error_and_debug(self):
+        self.shell_mock.side_effect = EnvironmentError("error")
+
+        self.assertRaises(
+            EnvironmentError, self.run_command, ["pull", "--shell-after", "--debug"]
+        )
+
+        self.pack_project_mock.assert_not_called()
+        self.execute_step_mock.assert_called_once_with(steps.PULL)
+        self.shell_mock.assert_called_once_with()