containerize: pin the Spack version used in a container (#21910)

This PR permits to specify the `url` and `ref` of the Spack instance used in a container recipe simply by expanding the YAML schema as outlined in #20442:
```yaml
container:
  images:
    os: amazonlinux:2
    spack:
      ref: develop
      resolve_sha: true
```
The `resolve_sha` option, if true, verifies the `ref` by cloning the Spack repository in a temporary directory and transforming any tag or branch name to a commit sha. When this new ability is leveraged an additional "bootstrap" stage is added, which builds an image with Spack setup and ready to install software. The Spack repository to be used can be customized with the `url` keyword under `spack`.

Modifications:
- [x] Permit to pin the version of Spack, either by branch or tag or sha
- [x] Added a few new OSes (centos:8, amazonlinux:2, ubuntu:20.04, alpine:3, cuda:11.2.1)
- [x] Permit to print the bootstrap image as a standalone
- [x] Add documentation on the new part of the schema
- [x] Add unit tests for different use cases

11 files changed, 447 insertions, 86 deletions

diff --git a/lib/spack/docs/containers.rst b/lib/spack/docs/containers.rst
index 4364b5d4db..3d32de0841 100644
--- a/lib/spack/docs/containers.rst
+++ b/lib/spack/docs/containers.rst
@@ -197,7 +197,7 @@ Setting Base Images
 
 The ``images`` subsection is used to select both the image where
 Spack builds the software and the image where the built software
-is installed. This attribute can be set in two different ways and
+is installed. This attribute can be set in different ways and
 which one to use depends on the use case at hand.
 
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -257,10 +257,54 @@ software is respectively built and installed:
 
    ENTRYPOINT ["/bin/bash", "--rcfile", "/etc/profile", "-l"]
 
-This method of selecting base images is the simplest of the two, and we advise
+This is the simplest available method of selecting base images, and we advise
 to use it whenever possible. There are cases though where using Spack official
-images is not enough to fit production needs. In these situations users can manually
-select which base image to start from in the recipe, as we'll see next.
+images is not enough to fit production needs. In these situations users can
+extend the recipe to start with the bootstrapping of Spack at a certain pinned
+version or manually select which base image to start from in the recipe,
+as we'll see next.
+
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+Use a Bootstrap Stage for Spack
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+In some cases users may want to pin the commit sha that is used for Spack, to ensure later
+reproducibility, or start from a fork of the official Spack repository to try a bugfix or
+a feature in the early stage of development. This is possible by being just a little more
+verbose when specifying information about Spack in the ``spack.yaml`` file:
+
+.. code-block:: yaml
+
+   images:
+     os: amazonlinux:2
+     spack:
+       # URL of the Spack repository to be used in the container image
+       url: <to-use-a-fork>
+       # Either a commit sha, a branch name or a tag
+       ref: <sha/tag/branch>
+       # If true turn a branch name or a tag into the corresponding commit
+       # sha at the time of recipe generation
+       resolve_sha: <true/false>
+
+``url`` specifies the URL from which to clone Spack and defaults to https://github.com/spack/spack.
+The ``ref`` attribute can be either a commit sha, a branch name or a tag. The default value in
+this case is to use the ``develop`` branch, but it may change in the future to point to the latest stable
+release. Finally ``resolve_sha`` transform branch names or tags into the corresponding commit
+shas at the time of recipe generation, to allow for a greater reproducibility of the results
+at a later time.
+
+The list of operating systems that can be used to bootstrap Spack can be
+obtained with:
+
+.. command-output:: spack containerize --list-os
+
+.. note::
+
+   The ``resolve_sha`` option uses ``git rev-parse`` under the hood and thus it requires
+   to checkout the corresponding Spack repository in a temporary folder before generating
+   the recipe. Recipe generation may take longer when this option is set to true because
+   of this additional step.
+
 
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 Use Custom Images Provided by Users
@@ -412,6 +456,18 @@ to customize the generation of container recipes:
      - Version of Spack use in the ``build`` stage
      - Valid tags for ``base:image``
      - Yes, if using constrained selection of base images
+   * - ``images:spack:url``
+     - Repository from which Spack is cloned
+     - Any fork of Spack
+     - No
+   * - ``images:spack:ref``
+     - Reference for the checkout of Spack
+     - Either a commit sha, a branch name or a tag
+     - No
+   * - ``images:spack:resolve_sha``
+     - Resolve branches and tags in ``spack.yaml`` to commits in the generated recipe
+     - True or False (default: False)
+     - No
    * - ``images:build``
      - Image to be used in the ``build`` stage
      - Any valid container image
diff --git a/lib/spack/llnl/util/filesystem.py b/lib/spack/llnl/util/filesystem.py
index f3c2ee5ab1..4c4ea1d5b8 100644
--- a/lib/spack/llnl/util/filesystem.py
+++ b/lib/spack/llnl/util/filesystem.py
@@ -1855,3 +1855,18 @@ def keep_modification_time(*filenames):
     for f, mtime in mtimes.items():
         if os.path.exists(f):
             os.utime(f, (os.path.getatime(f), mtime))
+
+
+@contextmanager
+def temporary_dir(*args, **kwargs):
+    """Create a temporary directory and cd's into it. Delete the directory
+    on exit.
+
+    Takes the same arguments as tempfile.mkdtemp()
+    """
+    tmp_dir = tempfile.mkdtemp(*args, **kwargs)
+    try:
+        with working_dir(tmp_dir):
+            yield tmp_dir
+    finally:
+        remove_directory_contents(tmp_dir)
diff --git a/lib/spack/spack/cmd/containerize.py b/lib/spack/spack/cmd/containerize.py
index ea90b24b87..e22a5b4c4e 100644
--- a/lib/spack/spack/cmd/containerize.py
+++ b/lib/spack/spack/cmd/containerize.py
@@ -5,7 +5,10 @@
 import os
 import os.path
 
+import llnl.util.tty
+
 import spack.container
+import spack.container.images
 import spack.monitor
 
 description = ("creates recipes to build images for different"
@@ -16,9 +19,26 @@ level = "long"
 
 def setup_parser(subparser):
     monitor_group = spack.monitor.get_monitor_group(subparser)  # noqa
+    subparser.add_argument(
+        '--list-os', action='store_true', default=False,
+        help='list all the OS that can be used in the bootstrap phase and exit'
+    )
+    subparser.add_argument(
+        '--last-stage',
+        choices=('bootstrap', 'build', 'final'),
+        default='final',
+        help='last stage in the container recipe'
+    )
 
 
 def containerize(parser, args):
+    if args.list_os:
+        possible_os = spack.container.images.all_bootstrap_os()
+        msg = 'The following operating systems can be used to bootstrap Spack:'
+        msg += '\n{0}'.format(' '.join(possible_os))
+        llnl.util.tty.msg(msg)
+        return
+
     config_dir = args.env_dir or os.getcwd()
     config_file = os.path.abspath(os.path.join(config_dir, 'spack.yaml'))
     if not os.path.exists(config_file):
@@ -29,10 +49,12 @@ def containerize(parser, args):
 
     # If we have a monitor request, add monitor metadata to config
     if args.use_monitor:
-        config['spack']['monitor'] = {"disable_auth": args.monitor_disable_auth,
-                                      "host": args.monitor_host,
-                                      "keep_going": args.monitor_keep_going,
-                                      "prefix": args.monitor_prefix,
-                                      "tags": args.monitor_tags}
-    recipe = spack.container.recipe(config)
+        config['spack']['monitor'] = {
+            "disable_auth": args.monitor_disable_auth,
+            "host": args.monitor_host,
+            "keep_going": args.monitor_keep_going,
+            "prefix": args.monitor_prefix,
+            "tags": args.monitor_tags
+        }
+    recipe = spack.container.recipe(config, last_phase=args.last_stage)
     print(recipe)
diff --git a/lib/spack/spack/container/images.json b/lib/spack/spack/container/images.json
index 9461d576d1..ee4e5a2caa 100644
--- a/lib/spack/spack/container/images.json
+++ b/lib/spack/spack/container/images.json
@@ -1,6 +1,53 @@
 {
   "images": {
+    "alpine:3": {
+      "bootstrap": {
+        "template": "container/alpine_3.dockerfile"
+      },
+      "os_package_manager": "apk"
+    },
+    "amazonlinux:2": {
+      "bootstrap": {
+        "template": "container/amazonlinux_2.dockerfile"
+      },
+      "os_package_manager": "yum_amazon"
+    },
+    "centos:8": {
+      "bootstrap": {
+        "template": "container/centos_8.dockerfile"
+      },
+      "os_package_manager": "yum"
+    },
+    "centos:7": {
+      "bootstrap": {
+        "template": "container/centos_7.dockerfile"
+      },
+      "os_package_manager": "yum",
+      "build": "spack/centos7",
+      "build_tags": {
+        "develop": "latest"
+      }
+    },
+    "nvidia/cuda:11.2.1": {
+      "bootstrap": {
+        "template": "container/cuda_11_2_1.dockerfile",
+        "image": "nvidia/cuda:11.2.1-devel"
+      },
+      "final": {
+        "image": "nvidia/cuda:11.2.1-base"
+      },
+      "os_package_manager": "apt"
+    },
+    "ubuntu:20.04": {
+      "bootstrap": {
+        "template": "container/ubuntu_2004.dockerfile"
+      },
+      "os_package_manager": "apt"
+    },
     "ubuntu:18.04": {
+      "bootstrap": {
+        "template": "container/ubuntu_1804.dockerfile"
+      },
       "os_package_manager": "apt",
       "build": "spack/ubuntu-bionic",
       "build_tags": {
@@ -8,22 +55,22 @@
       }
     },
     "ubuntu:16.04": {
+      "bootstrap": {
+        "template": "container/ubuntu_1604.dockerfile"
+      },
       "os_package_manager": "apt",
       "build": "spack/ubuntu-xenial",
       "build_tags": {
         "develop": "latest"
       }
-    },
-    "centos:7": {
-      "os_package_manager": "yum",
-      "environment": [],
-      "build": "spack/centos7",
-      "build_tags": {
-        "develop": "latest"
-      }
     }
   },
   "os_package_managers": {
+    "apk": {
+      "update": "apk update",
+      "install": "apk add --no-cache",
+      "clean": "true"
+    },
     "apt": {
       "update": "apt-get -yqq update && apt-get -yqq upgrade",
       "install": "apt-get -yqq install",
@@ -33,6 +80,11 @@
       "update": "yum update -y && yum install -y epel-release && yum update -y",
       "install": "yum install -y",
       "clean": "rm -rf /var/cache/yum  && yum clean all"
+    },
+    "yum_amazon": {
+      "update": "yum update -y && amazon-linux-extras install epel -y",
+      "install": "yum install -y",
+      "clean": "rm -rf /var/cache/yum  && yum clean all"
     }
   }
 }
diff --git a/lib/spack/spack/container/images.py b/lib/spack/spack/container/images.py
index 9d2e15f195..03591e68ee 100644
--- a/lib/spack/spack/container/images.py
+++ b/lib/spack/spack/container/images.py
@@ -2,9 +2,15 @@
 # Spack Project Developers. See the top-level COPYRIGHT file for details.
 #
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
-"""Manages the details on the images used in the build and the run stage."""
+"""Manages the details on the images used in the various stages."""
 import json
 import os.path
+import sys
+
+import llnl.util.filesystem as fs
+import llnl.util.tty as tty
+
+import spack.util.executable as executable
 
 #: Global variable used to cache in memory the content of images.json
 _data = None
@@ -39,18 +45,12 @@ def build_info(image, spack_version):
     # Don't handle error here, as a wrong image should have been
     # caught by the JSON schema
     image_data = data()["images"][image]
-    build_image = image_data['build']
-
-    # Try to check if we have a tag for this Spack version
-    try:
-        # Translate version from git to docker if necessary
-        build_tag = image_data['build_tags'].get(spack_version, spack_version)
-    except KeyError:
-        msg = ('the image "{0}" has no tag for Spack version "{1}" '
-               '[valid versions are {2}]')
-        msg = msg.format(build_image, spack_version,
-                         ', '.join(image_data['build_tags'].keys()))
-        raise ValueError(msg)
+    build_image = image_data.get('build', None)
+    if not build_image:
+        return None, None
+
+    # Translate version from git to docker if necessary
+    build_tag = image_data['build_tags'].get(spack_version, spack_version)
 
     return build_image, build_tag
 
@@ -70,6 +70,11 @@ def os_package_manager_for(image):
     return name
 
 
+def all_bootstrap_os():
+    """Return a list of all the OS that can be used to bootstrap Spack"""
+    return list(data()['images'])
+
+
 def commands_for(package_manager):
     """Returns the commands used to update system repositories, install
     system packages and clean afterwards.
@@ -82,3 +87,47 @@ def commands_for(package_manager):
     """
     info = data()["os_package_managers"][package_manager]
     return info['update'], info['install'], info['clean']
+
+
+def bootstrap_template_for(image):
+    return data()["images"][image]["bootstrap"]["template"]
+
+
+def _verify_ref(url, ref, enforce_sha):
+    # Do a checkout in a temporary directory
+    msg = 'Cloning "{0}" to verify ref "{1}"'.format(url, ref)
+    tty.info(msg, stream=sys.stderr)
+    git = executable.which('git', required=True)
+    with fs.temporary_dir():
+        git('clone', '-q', url, '.')
+        sha = git('rev-parse', '-q', ref + '^{commit}',
+                  output=str, error=os.devnull, fail_on_error=False)
+        if git.returncode:
+            msg = '"{0}" is not a valid reference for "{1}"'
+            raise RuntimeError(msg.format(sha, url))
+
+        if enforce_sha:
+            ref = sha.strip()
+
+        return ref
+
+
+def checkout_command(url, ref, enforce_sha, verify):
+    """Return the checkout command to be used in the bootstrap phase.
+
+    Args:
+        url (str): url of the Spack repository
+        ref (str): either a branch name, a tag or a commit sha
+        enforce_sha (bool): if true turns every
+        verify (bool):
+    """
+    url = url or 'https://github.com/spack/spack.git'
+    ref = ref or 'develop'
+    enforce_sha, verify = bool(enforce_sha), bool(verify)
+    # If we want to enforce a sha or verify the ref we need
+    # to checkout the repository locally
+    if enforce_sha or verify:
+        ref = _verify_ref(url, ref, enforce_sha)
+
+    command = 'git clone {0} . && git checkout {1} '.format(url, ref)
+    return command
diff --git a/lib/spack/spack/container/writers/__init__.py b/lib/spack/spack/container/writers/__init__.py
index 978e999605..9808969bfc 100644
--- a/lib/spack/spack/container/writers/__init__.py
+++ b/lib/spack/spack/container/writers/__init__.py
@@ -12,7 +12,14 @@ import spack.environment as ev
 import spack.schema.env
 import spack.tengine as tengine
 import spack.util.spack_yaml as syaml
-from spack.container.images import build_info, commands_for, os_package_manager_for
+from spack.container.images import (
+    bootstrap_template_for,
+    build_info,
+    checkout_command,
+    commands_for,
+    data,
+    os_package_manager_for,
+)
 
 #: Caches all the writers that are currently supported
 _writer_factory = {}
@@ -31,23 +38,94 @@ def writer(name):
     return _decorator
 
 
-def create(configuration):
+def create(configuration, last_phase=None):
     """Returns a writer that conforms to the configuration passed as input.
 
     Args:
-        configuration: how to generate the current recipe
+        configuration (dict): how to generate the current recipe
+        last_phase (str): last phase to be printed or None to print them all
     """
     name = ev.config_dict(configuration)['container']['format']
-    return _writer_factory[name](configuration)
+    return _writer_factory[name](configuration, last_phase)
 
 
-def recipe(configuration):
+def recipe(configuration, last_phase=None):
     """Returns a recipe that conforms to the configuration passed as input.
 
     Args:
-        configuration: how to generate the current recipe
+        configuration (dict): how to generate the current recipe
+        last_phase (str): last phase to be printed or None to print them all
     """
-    return create(configuration)()
+    return create(configuration, last_phase)()
+
+
+def _stage_base_images(images_config):
+    """Return a tuple with the base images to be used at the various stages.
+
+    Args:
+        images_config (dict): configuration under container:images
+    """
+    # If we have custom base images, just return them verbatim.
+    build_stage = images_config.get('build', None)
+    if build_stage:
+        final_stage = images_config['final']
+        return None, build_stage, final_stage
+
+    # Check the operating system: this will be the base of the bootstrap
+    # stage, if there, and of the final stage.
+    operating_system = images_config.get('os', None)
+
+    # Check the OS is mentioned in the internal data stored in a JSON file
+    images_json = data()['images']
+    if not any(os_name == operating_system for os_name in images_json):
+        msg = ('invalid operating system name "{0}". '
+               '[Allowed values are {1}]')
+        msg = msg.format(operating_system, ', '.join(data()['images']))
+        raise ValueError(msg)
+
+    # Retrieve the build stage
+    spack_info = images_config['spack']
+    if isinstance(spack_info, dict):
+        build_stage = 'bootstrap'
+    else:
+        spack_version = images_config['spack']
+        image_name, tag = build_info(operating_system, spack_version)
+        build_stage = 'bootstrap'
+        if image_name:
+            build_stage = ':'.join([image_name, tag])
+
+    # Retrieve the bootstrap stage
+    bootstrap_stage = None
+    if build_stage == 'bootstrap':
+        bootstrap_stage = images_json[operating_system]['bootstrap'].get(
+            'image', operating_system
+        )
+
+    # Retrieve the final stage
+    final_stage = images_json[operating_system].get(
+        'final', {'image': operating_system}
+    )['image']
+
+    return bootstrap_stage, build_stage, final_stage
+
+
+def _spack_checkout_config(images_config):
+    spack_info = images_config['spack']
+
+    url = 'https://github.com/spack/spack.git'
+    ref = 'develop'
+    resolve_sha, verify = False, False
+
+    # Config specific values may override defaults
+    if isinstance(spack_info, dict):
+        url = spack_info.get('url', url)
+        ref = spack_info.get('ref', ref)
+        resolve_sha = spack_info.get('resolve_sha', resolve_sha)
+        verify = spack_info.get('verify', verify)
+    else:
+        ref = spack_info
+
+    return url, ref, resolve_sha, verify
 
 
 class PathContext(tengine.Context):
@@ -55,41 +133,34 @@ class PathContext(tengine.Context):
     install software in a common location and make it available
     directly via PATH.
     """
-    def __init__(self, config):
+    def __init__(self, config, last_phase):
         self.config = ev.config_dict(config)
         self.container_config = self.config['container']
 
+        # Operating system tag as written in the configuration file
+        self.operating_system_key = self.container_config['images'].get('os')
+        # Get base images and verify the OS
+        bootstrap, build, final = _stage_base_images(
+            self.container_config['images']
+        )
+        self.bootstrap_image = bootstrap
+        self.build_image = build
+        self.final_image = final
+
+        # Record the last phase
+        self.last_phase = last_phase
+
     @tengine.context_property
     def run(self):
         """Information related to the run image."""
-        images_config = self.container_config['images']
-
-        # Check if we have custom images
-        image = images_config.get('final', None)
-        # If not use the base OS image
-        if image is None:
-            image = images_config['os']
-
         Run = collections.namedtuple('Run', ['image'])
-        return Run(image=image)
+        return Run(image=self.final_image)
 
     @tengine.context_property
     def build(self):
         """Information related to the build image."""
-        images_config = self.container_config['images']
-
-        # Check if we have custom images
-        image = images_config.get('build', None)
-
-        # If not select the correct build image based on OS and Spack version
-        if image is None:
-            operating_system = images_config['os']
-            spack_version = images_config['spack']
-            image_name, tag = build_info(operating_system, spack_version)
-            image = ':'.join([image_name, tag])
-
         Build = collections.namedtuple('Build', ['image'])
-        return Build(image=image)
+        return Build(image=self.build_image)
 
     @tengine.context_property
     def strip(self):
@@ -213,6 +284,39 @@ class PathContext(tengine.Context):
     def labels(self):
         return self.container_config.get('labels', {})
 
+    @tengine.context_property
+    def bootstrap(self):
+        """Information related to the build image."""
+        images_config = self.container_config['images']
+        bootstrap_recipe = None
+        if self.bootstrap_image:
+            config_args = _spack_checkout_config(images_config)
+            command = checkout_command(*config_args)
+            template_path = bootstrap_template_for(self.operating_system_key)
+            env = tengine.make_environment()
+            context = {"bootstrap": {
+                "image": self.bootstrap_image,
+                "spack_checkout": command
+            }}
+            bootstrap_recipe = env.get_template(template_path).render(**context)
+
+        Bootstrap = collections.namedtuple('Bootstrap', ['image', 'recipe'])
+        return Bootstrap(image=self.bootstrap_image, recipe=bootstrap_recipe)
+
+    @tengine.context_property
+    def render_phase(self):
+        render_bootstrap = bool(self.bootstrap_image)
+        render_build = not (self.last_phase == 'bootstrap')
+        render_final = self.last_phase in (None, 'final')
+        Render = collections.namedtuple(
+            'Render', ['bootstrap', 'build', 'final']
+        )
+        return Render(
+            bootstrap=render_bootstrap,
+            build=render_build,
+            final=render_final
+        )
+
     def __call__(self):
         """Returns the recipe as a string"""
         env = tengine.make_environment()
diff --git a/lib/spack/spack/schema/container.py b/lib/spack/spack/schema/container.py
index 2011b42853..411137dc77 100644
--- a/lib/spack/spack/schema/container.py
+++ b/lib/spack/spack/schema/container.py
@@ -8,15 +8,18 @@ _stages_from_dockerhub = {
     'type': 'object',
     'additionalProperties': False,
     'properties': {
-        'os': {
-            'type': 'string',
-            'enum': ['ubuntu:18.04',
-                     'ubuntu:16.04',
-                     'centos:7']
-        },
-        'spack': {
-            'type': 'string',
-        },
+        'os': {'type': 'string'},
+        'spack': {'anyOf': [
+            {'type': 'string'},
+            {'type': 'object',
+             'additional_properties': False,
+             'properties': {
+                 'url': {'type': 'string'},
+                 'ref': {'type': 'string'},
+                 'resolve_sha': {'type': 'boolean', 'default': False},
+                 'verify': {'type': 'boolean', 'default': False}
+             }}
+        ]},
     },
     'required': ['os', 'spack']
 }
diff --git a/lib/spack/spack/test/container/cli.py b/lib/spack/spack/test/container/cli.py
index 1dd0840d7f..ddbb9eca9b 100644
--- a/lib/spack/spack/test/container/cli.py
+++ b/lib/spack/spack/test/container/cli.py
@@ -2,8 +2,11 @@
 # Spack Project Developers. See the top-level COPYRIGHT file for details.
 #
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
+import pytest
+
 import llnl.util.filesystem as fs
 
+import spack.container.images
 import spack.main
 
 containerize = spack.main.SpackCommand('containerize')
@@ -14,3 +17,29 @@ def test_command(default_config, container_config_dir, capsys):
         with fs.working_dir(container_config_dir):
             output = containerize()
     assert 'FROM spack/ubuntu-bionic' in output
+
+
+def test_listing_possible_os():
+    output = containerize('--list-os')
+
+    for expected_os in spack.container.images.all_bootstrap_os():
+        assert expected_os in output
+
+
+@pytest.mark.maybeslow
+@pytest.mark.requires_executables('git')
+def test_bootstrap_phase(minimal_configuration, config_dumper, capsys):
+    minimal_configuration['spack']['container']['images'] = {
+        'os': 'amazonlinux:2',
+        'spack': {
+            'resolve_sha': True
+        }
+    }
+    spack_yaml_dir = config_dumper(minimal_configuration)
+
+    with capsys.disabled():
+        with fs.working_dir(spack_yaml_dir):
+            output = containerize()
+
+    # Check for the presence of the clone command
+    assert 'git clone' in output
diff --git a/lib/spack/spack/test/container/docker.py b/lib/spack/spack/test/container/docker.py
index 43a32b0720..0da4313731 100644
--- a/lib/spack/spack/test/container/docker.py
+++ b/lib/spack/spack/test/container/docker.py
@@ -2,6 +2,7 @@
 # Spack Project Developers. See the top-level COPYRIGHT file for details.
 #
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
+import pytest
 
 import spack.container.writers as writers
 
@@ -74,3 +75,42 @@ def test_extra_instructions_is_set_from_config(minimal_configuration):
     del e['extra_instructions']['build']
     writer = writers.create(minimal_configuration)
     assert writer.extra_instructions == (None, test_line)
+
+
+def test_custom_base_images(minimal_configuration):
+    """Test setting custom base images from configuration file"""
+    minimal_configuration['spack']['container']['images'] = {
+        'build': 'custom-build:latest',
+        'final': 'custom-final:latest'
+    }
+    writer = writers.create(minimal_configuration)
+
+    assert writer.bootstrap.image is None
+    assert writer.build.image == 'custom-build:latest'
+    assert writer.run.image == 'custom-final:latest'
+
+
+@pytest.mark.parametrize('images_cfg,expected', [
+    ({'os': 'amazonlinux:2', 'spack': 'develop'}, {
+        'bootstrap_image': 'amazonlinux:2',
+        'build_image': 'bootstrap',
+        'final_image': 'amazonlinux:2'
+    })
+])
+def test_base_images_with_bootstrap(
+        minimal_configuration, images_cfg, expected
+):
+    """Check that base images are computed correctly when a
+    bootstrap phase is present
+    """
+    minimal_configuration['spack']['container']['images'] = images_cfg
+    writer = writers.create(minimal_configuration)
+
+    for property_name, value in expected.items():
+        assert getattr(writer, property_name) == value
+
+
+def test_error_message_invalid_os(minimal_configuration):
+    minimal_configuration['spack']['container']['images']['os'] = 'invalid:1'
+    with pytest.raises(ValueError, match='invalid operating system'):
+        writers.create(minimal_configuration)
diff --git a/lib/spack/spack/test/container/schema.py b/lib/spack/spack/test/container/schema.py
deleted file mode 100644
index 4bb0d574a9..0000000000
--- a/lib/spack/spack/test/container/schema.py
+++ /dev/null
@@ -1,16 +0,0 @@
-# Copyright 2013-2021 Lawrence Livermore National Security, LLC and other
-# Spack Project Developers. See the top-level COPYRIGHT file for details.
-#
-# SPDX-License-Identifier: (Apache-2.0 OR MIT)
-
-import spack.container
-import spack.schema.container
-
-
-def test_images_in_schema():
-    properties = spack.schema.container.container_schema['properties']
-    allowed_images = set(
-        properties['images']['anyOf'][0]['properties']['os']['enum']
-    )
-    images_in_json = set(x for x in spack.container.images.data()['images'])
-    assert images_in_json == allowed_images
diff --git a/lib/spack/spack/test/llnl/util/filesystem.py b/lib/spack/spack/test/llnl/util/filesystem.py
index 0db8e93ae8..39b81d574f 100644
--- a/lib/spack/spack/test/llnl/util/filesystem.py
+++ b/lib/spack/spack/test/llnl/util/filesystem.py
@@ -610,3 +610,10 @@ def test_keep_modification_time(tmpdir):
     assert file1.read().strip() == 'file1'
     assert not file2.exists()
     assert int(mtime1) == int(file1.mtime())
+
+
+def test_temporary_dir_context_manager():
+    previous_dir = os.path.realpath(os.getcwd())
+    with fs.temporary_dir() as tmp_dir:
+        assert previous_dir != os.path.realpath(os.getcwd())
+        assert os.path.realpath(str(tmp_dir)) == os.path.realpath(os.getcwd())