diff options
Diffstat (limited to 'tests/testutils/artifactshare.py')
-rw-r--r-- | tests/testutils/artifactshare.py | 16 |
1 files changed, 15 insertions, 1 deletions
diff --git a/tests/testutils/artifactshare.py b/tests/testutils/artifactshare.py index 8d0448f8a..19c19131a 100644 --- a/tests/testutils/artifactshare.py +++ b/tests/testutils/artifactshare.py @@ -13,7 +13,7 @@ from buildstream._cas import CASCache from buildstream._cas.casserver import create_server from buildstream._exceptions import CASError from buildstream._protos.build.bazel.remote.execution.v2 import remote_execution_pb2 -from buildstream._protos.buildstream.v2 import artifact_pb2 +from buildstream._protos.buildstream.v2 import artifact_pb2, source_pb2 class BaseArtifactShare: @@ -120,6 +120,8 @@ class ArtifactShare(BaseArtifactShare): os.makedirs(self.repodir) self.artifactdir = os.path.join(self.repodir, "artifacts", "refs") os.makedirs(self.artifactdir) + self.sourcedir = os.path.join(self.repodir, "source_protos", "refs") + os.makedirs(self.sourcedir) self.cas = CASCache(self.repodir, casd=casd) @@ -160,6 +162,18 @@ class ArtifactShare(BaseArtifactShare): return artifact_proto + def get_source_proto(self, source_name): + source_proto = source_pb2.Source() + source_path = os.path.join(self.sourcedir, source_name) + + try: + with open(source_path, "rb") as f: + source_proto.ParseFromString(f.read()) + except FileNotFoundError: + return None + + return source_proto + def get_cas_files(self, artifact_proto): reachable = set() |