_artifact.py: Add cached and strong_key parameters to set_cached()

1 files changed, 14 insertions, 4 deletions

diff --git a/src/buildstream/_artifact.py b/src/buildstream/_artifact.py
index b63cff6e5..9809ac822 100644
--- a/src/buildstream/_artifact.py
+++ b/src/buildstream/_artifact.py
@@ -613,10 +613,20 @@ class Artifact:
     # Mark the artifact as cached without querying the filesystem.
     # This is used as optimization when we know the artifact is available.
     #
-    def set_cached(self):
-        self._proto = self._load_proto()
-        assert self._proto
-        self._cached = True
+    def set_cached(self, cached=True, *, strong_key=None):
+        if strong_key:
+            assert not self._cache_key or self._cache_key == strong_key
+            self._cache_key = strong_key
+
+        if cached:
+            self._proto = self._load_proto()
+            assert self._proto
+            if self._cache_key:
+                assert self._cache_key == self._proto.strong_key
+            else:
+                self._cache_key = self._proto.strong_key
+
+        self._cached = cached
 
     # pull()
     #