Merge branch 'juerg/virtual-artifact-directory' into 'master'

Use virtual artifact directory to stage and extract metadata

See merge request BuildStream/buildstream!1184

8 files changed, 139 insertions, 217 deletions

diff --git a/buildstream/_artifactcache.py b/buildstream/_artifactcache.py
index b73304fac..b317296ec 100644
--- a/buildstream/_artifactcache.py
+++ b/buildstream/_artifactcache.py
@@ -54,7 +54,6 @@ class ArtifactCacheSpec(CASRemoteSpec):
 class ArtifactCache():
     def __init__(self, context):
         self.context = context
-        self.extractdir = context.extractdir
 
         self.cas = context.get_cascache()
         self.casquota = context.get_casquota()
@@ -73,8 +72,6 @@ class ArtifactCache():
         self._has_fetch_remotes = False
         self._has_push_remotes = False
 
-        os.makedirs(self.extractdir, exist_ok=True)
-
     # setup_remotes():
     #
     # Sets up which remotes to use
@@ -423,51 +420,28 @@ class ArtifactCache():
     #    (int): The amount of space recovered in the cache, in bytes
     #
     def remove(self, ref):
-
-        # Remove extract if not used by other ref
-        tree = self.cas.resolve_ref(ref)
-        ref_name, ref_hash = os.path.split(ref)
-        extract = os.path.join(self.extractdir, ref_name, tree.hash)
-        keys_file = os.path.join(extract, 'meta', 'keys.yaml')
-        if os.path.exists(keys_file):
-            keys_meta = _yaml.load(keys_file)
-            keys = [keys_meta['strong'], keys_meta['weak']]
-            remove_extract = True
-            for other_hash in keys:
-                if other_hash == ref_hash:
-                    continue
-                remove_extract = False
-                break
-
-            if remove_extract:
-                utils._force_rmtree(extract)
-
         return self.cas.remove(ref)
 
-    # extract():
+    # get_artifact_directory():
     #
-    # Extract cached artifact for the specified Element if it hasn't
-    # already been extracted.
+    # Get virtual directory for cached artifact of the specified Element.
     #
     # Assumes artifact has previously been fetched or committed.
     #
     # Args:
     #     element (Element): The Element to extract
     #     key (str): The cache key to use
-    #     subdir (str): Optional specific subdir to extract
     #
     # Raises:
     #     ArtifactError: In cases there was an OSError, or if the artifact
     #                    did not exist.
     #
-    # Returns: path to extracted artifact
+    # Returns: virtual directory object
     #
-    def extract(self, element, key, subdir=None):
+    def get_artifact_directory(self, element, key):
         ref = element.get_artifact_name(key)
-
-        path = os.path.join(self.extractdir, element._get_project().name, element.normal_name)
-
-        return self.cas.extract(ref, path, subdir=subdir)
+        digest = self.cas.resolve_ref(ref, update_mtime=True)
+        return CasBasedDirectory(self.cas, digest)
 
     # commit():
     #
@@ -609,11 +583,6 @@ class ArtifactCache():
 
                 if self.cas.pull(ref, remote, progress=progress, subdir=subdir, excluded_subdirs=excluded_subdirs):
                     element.info("Pulled artifact {} <- {}".format(display_key, remote.spec.url))
-                    if subdir:
-                        # Attempt to extract subdir into artifact extract dir if it already exists
-                        # without containing the subdir. If the respective artifact extract dir does not
-                        # exist a complete extraction will complete.
-                        self.extract(element, key, subdir)
                     # no need to pull from additional remotes
                     return True
                 else:
diff --git a/buildstream/_cas/cascache.py b/buildstream/_cas/cascache.py
index 2978b9bfe..40d7efe78 100644
--- a/buildstream/_cas/cascache.py
+++ b/buildstream/_cas/cascache.py
@@ -135,53 +135,6 @@ class CASCache():
         # True if subdir content is cached or if empty as expected
         return os.path.exists(objpath)
 
-    # extract():
-    #
-    # Extract cached directory for the specified ref if it hasn't
-    # already been extracted.
-    #
-    # Args:
-    #     ref (str): The ref whose directory to extract
-    #     path (str): The destination path
-    #     subdir (str): Optional specific dir to extract
-    #
-    # Raises:
-    #     CASCacheError: In cases there was an OSError, or if the ref did not exist.
-    #
-    # Returns: path to extracted directory
-    #
-    def extract(self, ref, path, subdir=None):
-        tree = self.resolve_ref(ref, update_mtime=True)
-
-        originaldest = dest = os.path.join(path, tree.hash)
-
-        # If artifact is already extracted, check if the optional subdir
-        # has also been extracted. If the artifact has not been extracted
-        # a full extraction would include the optional subdir
-        if os.path.isdir(dest):
-            if subdir:
-                if not os.path.isdir(os.path.join(dest, subdir)):
-                    dest = os.path.join(dest, subdir)
-                    tree = self._get_subdir(tree, subdir)
-                else:
-                    return dest
-            else:
-                return dest
-
-        with utils._tempdir(prefix='tmp', dir=self.tmpdir) as tmpdir:
-            checkoutdir = os.path.join(tmpdir, ref)
-            self.checkout(checkoutdir, tree, can_link=True)
-
-            try:
-                utils.move_atomic(checkoutdir, dest)
-            except utils.DirectoryExistsError:
-                # Another process beat us to rename
-                pass
-            except OSError as e:
-                raise CASCacheError("Failed to extract directory for ref '{}': {}".format(ref, e)) from e
-
-        return originaldest
-
     # checkout():
     #
     # Checkout the specified directory digest.
diff --git a/buildstream/_context.py b/buildstream/_context.py
index 75edac39d..8a9f485be 100644
--- a/buildstream/_context.py
+++ b/buildstream/_context.py
@@ -18,6 +18,7 @@
 #        Tristan Van Berkom <tristan.vanberkom@codethink.co.uk>
 
 import os
+import shutil
 import datetime
 from collections import deque
 from collections.abc import Mapping
@@ -70,9 +71,6 @@ class Context():
         # The directory for CAS
         self.casdir = None
 
-        # Extract directory
-        self.extractdir = None
-
         # The directory for temporary files
         self.tmpdir = None
 
@@ -218,7 +216,6 @@ class Context():
             setattr(self, directory, path)
 
         # add directories not set by users
-        self.extractdir = os.path.join(self.cachedir, 'extract')
         self.tmpdir = os.path.join(self.cachedir, 'tmp')
         self.casdir = os.path.join(self.cachedir, 'cas')
         self.builddir = os.path.join(self.cachedir, 'build')
@@ -230,6 +227,13 @@ class Context():
             os.rename(old_casdir, self.casdir)
             os.symlink(self.casdir, old_casdir)
 
+        # Cleanup old extract directories
+        old_extractdirs = [os.path.join(self.cachedir, 'artifacts', 'extract'),
+                           os.path.join(self.cachedir, 'extract')]
+        for old_extractdir in old_extractdirs:
+            if os.path.isdir(old_extractdir):
+                shutil.rmtree(old_extractdir, ignore_errors=True)
+
         # Load quota configuration
         # We need to find the first existing directory in the path of our
         # cachedir - the cachedir may not have been created yet.
diff --git a/buildstream/element.py b/buildstream/element.py
index 365931e27..fff95d0f5 100644
--- a/buildstream/element.py
+++ b/buildstream/element.py
@@ -662,9 +662,8 @@ class Element(Plugin):
         self.__assert_cached()
 
         with self.timed_activity("Staging {}/{}".format(self.name, self._get_brief_display_key())):
-            # Get the extracted artifact
-            artifact_base, _ = self.__extract()
-            artifact = os.path.join(artifact_base, 'files')
+            artifact_vdir, _ = self.__get_artifact_directory()
+            files_vdir = artifact_vdir.descend(['files'])
 
             # Hard link it into the staging area
             #
@@ -687,11 +686,11 @@ class Element(Plugin):
             else:
                 link_filter = split_filter
 
-            result = vstagedir.import_files(artifact, filter_callback=link_filter,
+            result = vstagedir.import_files(files_vdir, filter_callback=link_filter,
                                             report_written=True, can_link=True)
 
             if update_mtimes:
-                copy_result = vstagedir.import_files(artifact, filter_callback=copy_filter,
+                copy_result = vstagedir.import_files(files_vdir, filter_callback=copy_filter,
                                                      report_written=True, update_mtime=True)
                 result = result.combine(copy_result)
 
@@ -1478,9 +1477,9 @@ class Element(Plugin):
 
             # Check if we have a cached buildtree to use
             elif usebuildtree:
-                artifact_base, _ = self.__extract()
-                import_dir = os.path.join(artifact_base, 'buildtree')
-                if not os.listdir(import_dir):
+                artifact_vdir, _ = self.__get_artifact_directory()
+                import_dir = artifact_vdir.descend(['buildtree'])
+                if import_dir.is_empty():
                     detail = "Element type either does not expect a buildtree or it was explictily cached without one."
                     self.warn("WARNING: {} Artifact contains an empty buildtree".format(self.name), detail=detail)
             else:
@@ -2645,14 +2644,10 @@ class Element(Plugin):
     def __compute_splits(self, include=None, exclude=None, orphans=True):
         filter_func = self.__split_filter_func(include=include, exclude=exclude, orphans=orphans)
 
-        artifact_base, _ = self.__extract()
-        basedir = os.path.join(artifact_base, 'files')
+        artifact_vdir, _ = self.__get_artifact_directory()
+        files_vdir = artifact_vdir.descend(['files'])
 
-        # FIXME: Instead of listing the paths in an extracted artifact,
-        #        we should be using a manifest loaded from the artifact
-        #        metadata.
-        #
-        element_files = utils.list_relative_paths(basedir)
+        element_files = files_vdir.list_relative_paths()
 
         if not filter_func:
             # No splitting requested, just report complete artifact
@@ -2676,30 +2671,43 @@ class Element(Plugin):
             self.__whitelist_regex = re.compile(expression)
         return self.__whitelist_regex.match(os.path.join(os.sep, path))
 
-    # __extract():
+    # __get_extract_key():
+    #
+    # Get the key used to extract the artifact
     #
-    # Extract an artifact and return the directory
+    # Returns:
+    #    (str): The key
+    #
+    def __get_extract_key(self):
+
+        context = self._get_context()
+        key = self.__strict_cache_key
+
+        # Use weak cache key, if artifact is missing for strong cache key
+        # and the context allows use of weak cache keys
+        if not context.get_strict() and not self.__artifacts.contains(self, key):
+            key = self._get_cache_key(strength=_KeyStrength.WEAK)
+
+        return key
+
+    # __get_artifact_directory():
+    #
+    # Get a virtual directory for the artifact contents
     #
     # Args:
     #    key (str): The key for the artifact to extract,
     #               or None for the default key
     #
     # Returns:
-    #    (str): The path to the extracted artifact
+    #    (Directory): The virtual directory object
     #    (str): The chosen key
     #
-    def __extract(self, key=None):
+    def __get_artifact_directory(self, key=None):
 
         if key is None:
-            context = self._get_context()
-            key = self.__strict_cache_key
+            key = self.__get_extract_key()
 
-            # Use weak cache key, if artifact is missing for strong cache key
-            # and the context allows use of weak cache keys
-            if not context.get_strict() and not self.__artifacts.contains(self, key):
-                key = self._get_cache_key(strength=_KeyStrength.WEAK)
-
-        return (self.__artifacts.extract(self, key), key)
+        return (self.__artifacts.get_artifact_directory(self, key), key)
 
     # __get_artifact_metadata_keys():
     #
@@ -2715,7 +2723,7 @@ class Element(Plugin):
     def __get_artifact_metadata_keys(self, key=None):
 
         # Now extract it and possibly derive the key
-        artifact_base, key = self.__extract(key)
+        artifact_vdir, key = self.__get_artifact_directory(key)
 
         # Now try the cache, once we're sure about the key
         if key in self.__metadata_keys:
@@ -2723,8 +2731,8 @@ class Element(Plugin):
                     self.__metadata_keys[key]['weak'])
 
         # Parse the expensive yaml now and cache the result
-        meta_file = os.path.join(artifact_base, 'meta', 'keys.yaml')
-        meta = _yaml.load(meta_file)
+        meta_file = artifact_vdir._objpath(['meta', 'keys.yaml'])
+        meta = _yaml.load(meta_file, shortname='meta/keys.yaml')
         strong_key = meta['strong']
         weak_key = meta['weak']
 
@@ -2747,15 +2755,15 @@ class Element(Plugin):
     def __get_artifact_metadata_dependencies(self, key=None):
 
         # Extract it and possibly derive the key
-        artifact_base, key = self.__extract(key)
+        artifact_vdir, key = self.__get_artifact_directory(key)
 
         # Now try the cache, once we're sure about the key
         if key in self.__metadata_dependencies:
             return self.__metadata_dependencies[key]
 
         # Parse the expensive yaml now and cache the result
-        meta_file = os.path.join(artifact_base, 'meta', 'dependencies.yaml')
-        meta = _yaml.load(meta_file)
+        meta_file = artifact_vdir._objpath(['meta', 'dependencies.yaml'])
+        meta = _yaml.load(meta_file, shortname='meta/dependencies.yaml')
 
         # Cache it under both strong and weak keys
         strong_key, weak_key = self.__get_artifact_metadata_keys(key)
@@ -2776,15 +2784,15 @@ class Element(Plugin):
     def __get_artifact_metadata_workspaced(self, key=None):
 
         # Extract it and possibly derive the key
-        artifact_base, key = self.__extract(key)
+        artifact_vdir, key = self.__get_artifact_directory(key)
 
         # Now try the cache, once we're sure about the key
         if key in self.__metadata_workspaced:
             return self.__metadata_workspaced[key]
 
         # Parse the expensive yaml now and cache the result
-        meta_file = os.path.join(artifact_base, 'meta', 'workspaced.yaml')
-        meta = _yaml.load(meta_file)
+        meta_file = artifact_vdir._objpath(['meta', 'workspaced.yaml'])
+        meta = _yaml.load(meta_file, shortname='meta/workspaced.yaml')
         workspaced = meta['workspaced']
 
         # Cache it under both strong and weak keys
@@ -2806,15 +2814,15 @@ class Element(Plugin):
     def __get_artifact_metadata_workspaced_dependencies(self, key=None):
 
         # Extract it and possibly derive the key
-        artifact_base, key = self.__extract(key)
+        artifact_vdir, key = self.__get_artifact_directory(key)
 
         # Now try the cache, once we're sure about the key
         if key in self.__metadata_workspaced_dependencies:
             return self.__metadata_workspaced_dependencies[key]
 
         # Parse the expensive yaml now and cache the result
-        meta_file = os.path.join(artifact_base, 'meta', 'workspaced-dependencies.yaml')
-        meta = _yaml.load(meta_file)
+        meta_file = artifact_vdir._objpath(['meta', 'workspaced-dependencies.yaml'])
+        meta = _yaml.load(meta_file, shortname='meta/workspaced-dependencies.yaml')
         workspaced = meta['workspaced-dependencies']
 
         # Cache it under both strong and weak keys
@@ -2832,23 +2840,26 @@ class Element(Plugin):
         assert self.__dynamic_public is None
 
         # Load the public data from the artifact
-        artifact_base, _ = self.__extract()
-        metadir = os.path.join(artifact_base, 'meta')
-        self.__dynamic_public = _yaml.load(os.path.join(metadir, 'public.yaml'))
+        artifact_vdir, _ = self.__get_artifact_directory()
+        meta_file = artifact_vdir._objpath(['meta', 'public.yaml'])
+        self.__dynamic_public = _yaml.load(meta_file, shortname='meta/public.yaml')
 
     def __load_build_result(self, keystrength):
         self.__assert_cached(keystrength=keystrength)
         assert self.__build_result is None
 
-        artifact_base, _ = self.__extract(key=self.__weak_cache_key if keystrength is _KeyStrength.WEAK
-                                          else self.__strict_cache_key)
+        if keystrength is _KeyStrength.WEAK:
+            key = self.__weak_cache_key
+        else:
+            key = self.__strict_cache_key
+
+        artifact_vdir, _ = self.__get_artifact_directory(key)
 
-        metadir = os.path.join(artifact_base, 'meta')
-        result_path = os.path.join(metadir, 'build-result.yaml')
-        if not os.path.exists(result_path):
+        if not artifact_vdir._exists(['meta', 'build-result.yaml']):
             self.__build_result = (True, "succeeded", None)
             return
 
+        result_path = artifact_vdir._objpath(['meta', 'build-result.yaml'])
         data = _yaml.load(result_path)
         self.__build_result = (data["success"], data.get("description"), data.get("detail"))
 
diff --git a/buildstream/storage/_casbaseddirectory.py b/buildstream/storage/_casbaseddirectory.py
index 2c1f465c3..bd72e7c1c 100644
--- a/buildstream/storage/_casbaseddirectory.py
+++ b/buildstream/storage/_casbaseddirectory.py
@@ -591,3 +591,15 @@ class CasBasedDirectory(Directory):
         if not self.ref:
             self.ref = self.cas_cache.add_object(buffer=self.pb2_directory.SerializeToString())
         return self.ref
+
+    def _objpath(self, path):
+        subdir = self.descend(path[:-1])
+        entry = subdir.index[path[-1]]
+        return self.cas_cache.objpath(entry.pb_object.digest)
+
+    def _exists(self, path):
+        try:
+            subdir = self.descend(path[:-1])
+            return path[-1] in subdir.index
+        except VirtualDirectoryError:
+            return False
diff --git a/tests/artifactcache/expiry.py b/tests/artifactcache/expiry.py
index 8ece6295c..e959634cb 100644
--- a/tests/artifactcache/expiry.py
+++ b/tests/artifactcache/expiry.py
@@ -395,56 +395,6 @@ def test_invalid_cache_quota(cli, datafiles, tmpdir, quota, err_domain, err_reas
         res.assert_main_error(err_domain, err_reason)
 
 
-@pytest.mark.datafiles(DATA_DIR)
-def test_extract_expiry(cli, datafiles, tmpdir):
-    project = os.path.join(datafiles.dirname, datafiles.basename)
-    element_path = 'elements'
-
-    cli.configure({
-        'cache': {
-            'quota': 10000000,
-        }
-    })
-
-    create_element_size('target.bst', project, element_path, [], 6000000)
-    res = cli.run(project=project, args=['build', 'target.bst'])
-    res.assert_success()
-    assert cli.get_element_state(project, 'target.bst') == 'cached'
-
-    # Force creating extract
-    res = cli.run(project=project, args=['artifact', 'checkout', 'target.bst',
-                                         '--directory', os.path.join(str(tmpdir), 'checkout')])
-    res.assert_success()
-
-    # Get a snapshot of the extracts in advance
-    extractdir = os.path.join(project, 'cache', 'extract', 'test', 'target')
-    extracts = os.listdir(extractdir)
-    assert(len(extracts) == 1)
-    extract = os.path.join(extractdir, extracts[0])
-
-    # Remove target.bst from artifact cache
-    create_element_size('target2.bst', project, element_path, [], 6000000)
-    res = cli.run(project=project, args=['build', 'target2.bst'])
-    res.assert_success()
-    assert cli.get_element_state(project, 'target.bst') != 'cached'
-
-    # Now the extract should be removed.
-    assert not os.path.exists(extract)
-
-    # As an added bonus, let's ensure that no directories have been left behind
-    #
-    # Now we should have a directory for the cached target2.bst, which
-    # replaced target.bst in the cache, we should not have a directory
-    # for the target.bst
-    refsdir = os.path.join(project, 'cache', 'cas', 'refs', 'heads')
-    refsdirtest = os.path.join(refsdir, 'test')
-    refsdirtarget = os.path.join(refsdirtest, 'target')
-    refsdirtarget2 = os.path.join(refsdirtest, 'target2')
-
-    assert os.path.isdir(refsdirtarget2)
-    assert not os.path.exists(refsdirtarget)
-
-
 # Ensures that when launching BuildStream with a full artifact cache,
 # the cache size and cleanup jobs are run before any other jobs.
 #
diff --git a/tests/integration/artifact.py b/tests/integration/artifact.py
index 35cad2599..58e58bdcf 100644
--- a/tests/integration/artifact.py
+++ b/tests/integration/artifact.py
@@ -18,13 +18,17 @@
 #  Authors: Richard Maw <richard.maw@codethink.co.uk>
 #
 
+from contextlib import contextmanager
 import os
 import pytest
 import shutil
+import tempfile
 
+from buildstream import utils
 from buildstream.plugintestutils import cli_integration as cli
 from tests.testutils import create_artifact_share
 from tests.testutils.site import HAVE_SANDBOX
+from buildstream._cas import CASCache
 from buildstream._exceptions import ErrorDomain
 
 pytestmark = pytest.mark.integration
@@ -56,6 +60,16 @@ def test_cache_buildtrees(cli, tmpdir, datafiles):
             'cachedir': str(tmpdir)
         })
 
+        @contextmanager
+        def cas_extract_buildtree(digest):
+            extractdir = tempfile.mkdtemp(prefix="tmp", dir=str(tmpdir))
+            try:
+                cas = CASCache(str(tmpdir))
+                cas.checkout(extractdir, digest)
+                yield os.path.join(extractdir, 'buildtree')
+            finally:
+                utils._force_rmtree(extractdir)
+
         # Build autotools element with cache-buildtrees set via the
         # cli. The artifact should be successfully pushed to the share1 remote
         # and cached locally with an 'empty' buildtree digest, as it's not a
@@ -69,22 +83,20 @@ def test_cache_buildtrees(cli, tmpdir, datafiles):
         # to not cache buildtrees
         cache_key = cli.get_element_key(project, element_name)
         elementdigest = share1.has_artifact('test', element_name, cache_key)
-        buildtreedir = os.path.join(str(tmpdir), 'extract', 'test', 'autotools-amhello',
-                                    elementdigest.hash, 'buildtree')
-        assert os.path.isdir(buildtreedir)
-        assert not os.listdir(buildtreedir)
+        with cas_extract_buildtree(elementdigest) as buildtreedir:
+            assert os.path.isdir(buildtreedir)
+            assert not os.listdir(buildtreedir)
 
         # Delete the local cached artifacts, and assert the when pulled with --pull-buildtrees
         # that is was cached in share1 as expected with an empty buildtree dir
         shutil.rmtree(os.path.join(str(tmpdir), 'cas'))
-        shutil.rmtree(os.path.join(str(tmpdir), 'extract'))
         assert cli.get_element_state(project, element_name) != 'cached'
         result = cli.run(project=project, args=['--pull-buildtrees', 'artifact', 'pull', element_name])
         assert element_name in result.get_pulled_elements()
-        assert os.path.isdir(buildtreedir)
-        assert not os.listdir(buildtreedir)
+        with cas_extract_buildtree(elementdigest) as buildtreedir:
+            assert os.path.isdir(buildtreedir)
+            assert not os.listdir(buildtreedir)
         shutil.rmtree(os.path.join(str(tmpdir), 'cas'))
-        shutil.rmtree(os.path.join(str(tmpdir), 'extract'))
 
         # Assert that the default behaviour of pull to not include buildtrees on the artifact
         # in share1 which was purposely cached with an empty one behaves as expected. As such the
@@ -92,9 +104,9 @@ def test_cache_buildtrees(cli, tmpdir, datafiles):
         # leading to no buildtreedir being extracted
         result = cli.run(project=project, args=['artifact', 'pull', element_name])
         assert element_name in result.get_pulled_elements()
-        assert not os.path.isdir(buildtreedir)
+        with cas_extract_buildtree(elementdigest) as buildtreedir:
+            assert not os.path.isdir(buildtreedir)
         shutil.rmtree(os.path.join(str(tmpdir), 'cas'))
-        shutil.rmtree(os.path.join(str(tmpdir), 'extract'))
 
         # Repeat building the artifacts, this time with the default behaviour of caching buildtrees,
         # as such the buildtree dir should not be empty
@@ -109,22 +121,20 @@ def test_cache_buildtrees(cli, tmpdir, datafiles):
 
         # Cache key will be the same however the digest hash will have changed as expected, so reconstruct paths
         elementdigest = share2.has_artifact('test', element_name, cache_key)
-        buildtreedir = os.path.join(str(tmpdir), 'extract', 'test', 'autotools-amhello',
-                                    elementdigest.hash, 'buildtree')
-        assert os.path.isdir(buildtreedir)
-        assert os.listdir(buildtreedir) is not None
+        with cas_extract_buildtree(elementdigest) as buildtreedir:
+            assert os.path.isdir(buildtreedir)
+            assert os.listdir(buildtreedir) is not None
 
         # Delete the local cached artifacts, and assert that when pulled with --pull-buildtrees
         # that it was cached in share2 as expected with a populated buildtree dir
         shutil.rmtree(os.path.join(str(tmpdir), 'cas'))
-        shutil.rmtree(os.path.join(str(tmpdir), 'extract'))
         assert cli.get_element_state(project, element_name) != 'cached'
         result = cli.run(project=project, args=['--pull-buildtrees', 'artifact', 'pull', element_name])
         assert element_name in result.get_pulled_elements()
-        assert os.path.isdir(buildtreedir)
-        assert os.listdir(buildtreedir) is not None
+        with cas_extract_buildtree(elementdigest) as buildtreedir:
+            assert os.path.isdir(buildtreedir)
+            assert os.listdir(buildtreedir) is not None
         shutil.rmtree(os.path.join(str(tmpdir), 'cas'))
-        shutil.rmtree(os.path.join(str(tmpdir), 'extract'))
 
         # Clarify that the user config option for cache-buildtrees works as the cli
         # main option does. Point to share3 which does not have the artifacts cached to force
@@ -139,7 +149,6 @@ def test_cache_buildtrees(cli, tmpdir, datafiles):
         assert cli.get_element_state(project, element_name) == 'cached'
         cache_key = cli.get_element_key(project, element_name)
         elementdigest = share3.has_artifact('test', element_name, cache_key)
-        buildtreedir = os.path.join(str(tmpdir), 'extract', 'test', 'autotools-amhello',
-                                    elementdigest.hash, 'buildtree')
-        assert os.path.isdir(buildtreedir)
-        assert not os.listdir(buildtreedir)
+        with cas_extract_buildtree(elementdigest) as buildtreedir:
+            assert os.path.isdir(buildtreedir)
+            assert not os.listdir(buildtreedir)
diff --git a/tests/integration/pullbuildtrees.py b/tests/integration/pullbuildtrees.py
index 538ed8c37..1d01ca383 100644
--- a/tests/integration/pullbuildtrees.py
+++ b/tests/integration/pullbuildtrees.py
@@ -1,13 +1,17 @@
+from contextlib import contextmanager
 import os
-import shutil
 import pytest
+import shutil
+import tempfile
 
 from tests.testutils import create_artifact_share
 from tests.testutils.site import HAVE_SANDBOX
 
+from buildstream import utils
 from buildstream.plugintestutils.integration import assert_contains
 from buildstream.plugintestutils import cli, cli_integration as cli2
 from buildstream.plugintestutils.integration import assert_contains
+from buildstream._cas import CASCache
 from buildstream._exceptions import ErrorDomain, LoadErrorReason
 
 
@@ -22,7 +26,6 @@ DATA_DIR = os.path.join(
 # cleared as just forcefully removing the refpath leaves dangling objects.
 def default_state(cli, tmpdir, share):
     shutil.rmtree(os.path.join(str(tmpdir), 'cas'))
-    shutil.rmtree(os.path.join(str(tmpdir), 'extract'))
     cli.configure({
         'artifacts': {'url': share.repo, 'push': False},
         'cachedir': str(tmpdir),
@@ -49,6 +52,16 @@ def test_pullbuildtrees(cli2, tmpdir, datafiles):
             'cachedir': str(tmpdir),
         })
 
+        @contextmanager
+        def cas_extract_buildtree(digest):
+            extractdir = tempfile.mkdtemp(prefix="tmp", dir=str(tmpdir))
+            try:
+                cas = CASCache(str(tmpdir))
+                cas.checkout(extractdir, digest)
+                yield os.path.join(extractdir, 'buildtree')
+            finally:
+                utils._force_rmtree(extractdir)
+
         # Build autotools element, checked pushed, delete local
         result = cli2.run(project=project, args=['build', element_name])
         assert result.exit_code == 0
@@ -70,17 +83,16 @@ def test_pullbuildtrees(cli2, tmpdir, datafiles):
 
         # Pull artifact with default config, then assert that pulling
         # with buildtrees cli flag set creates a pull job.
-        # Also assert that the buildtree is added to the artifact's
-        # extract dir
+        # Also assert that the buildtree is added to the local CAS.
         result = cli2.run(project=project, args=['artifact', 'pull', element_name])
         assert element_name in result.get_pulled_elements()
         elementdigest = share1.has_artifact('test', element_name, cli2.get_element_key(project, element_name))
-        buildtreedir = os.path.join(str(tmpdir), 'extract', 'test', 'autotools-amhello',
-                                    elementdigest.hash, 'buildtree')
-        assert not os.path.isdir(buildtreedir)
+        with cas_extract_buildtree(elementdigest) as buildtreedir:
+            assert not os.path.isdir(buildtreedir)
         result = cli2.run(project=project, args=['--pull-buildtrees', 'artifact', 'pull', element_name])
         assert element_name in result.get_pulled_elements()
-        assert os.path.isdir(buildtreedir)
+        with cas_extract_buildtree(elementdigest) as buildtreedir:
+            assert os.path.isdir(buildtreedir)
         default_state(cli2, tmpdir, share1)
 
         # Pull artifact with pullbuildtrees set in user config, then assert
@@ -138,7 +150,8 @@ def test_pullbuildtrees(cli2, tmpdir, datafiles):
         result = cli2.run(project=project, args=['--pull-buildtrees', 'artifact', 'push', element_name])
         assert "Attempting to fetch missing artifact buildtrees" in result.stderr
         assert element_name not in result.get_pulled_elements()
-        assert not os.path.isdir(buildtreedir)
+        with cas_extract_buildtree(elementdigest) as buildtreedir:
+            assert not os.path.isdir(buildtreedir)
         assert element_name not in result.get_pushed_elements()
         assert not share3.has_artifact('test', element_name, cli2.get_element_key(project, element_name))
 
@@ -150,7 +163,8 @@ def test_pullbuildtrees(cli2, tmpdir, datafiles):
         result = cli2.run(project=project, args=['--pull-buildtrees', 'artifact', 'push', element_name])
         assert "Attempting to fetch missing artifact buildtrees" in result.stderr
         assert element_name in result.get_pulled_elements()
-        assert os.path.isdir(buildtreedir)
+        with cas_extract_buildtree(elementdigest) as buildtreedir:
+            assert os.path.isdir(buildtreedir)
         assert element_name in result.get_pushed_elements()
         assert share3.has_artifact('test', element_name, cli2.get_element_key(project, element_name))