diff options
Diffstat (limited to 'tests/testutils/artifactshare.py')
-rw-r--r-- | tests/testutils/artifactshare.py | 34 |
1 files changed, 23 insertions, 11 deletions
diff --git a/tests/testutils/artifactshare.py b/tests/testutils/artifactshare.py index 07def5c86..bd9c97c61 100644 --- a/tests/testutils/artifactshare.py +++ b/tests/testutils/artifactshare.py @@ -15,10 +15,11 @@ from buildstream._cas.casserver import create_server from buildstream._exceptions import CASError from buildstream._protos.build.bazel.remote.asset.v1 import remote_asset_pb2, remote_asset_pb2_grpc from buildstream._protos.build.bazel.remote.execution.v2 import remote_execution_pb2 -from buildstream._protos.buildstream.v2 import artifact_pb2, source_pb2 +from buildstream._protos.buildstream.v2 import artifact_pb2 from buildstream._protos.google.rpc import code_pb2 REMOTE_ASSET_ARTIFACT_URN_TEMPLATE = "urn:fdc:buildstream.build:2020:artifact:{}" +REMOTE_ASSET_SOURCE_URN_TEMPLATE = "urn:fdc:buildstream.build:2020:source:{}" class BaseArtifactShare: @@ -123,8 +124,6 @@ class ArtifactShare(BaseArtifactShare): # self.repodir = os.path.join(self.directory, "repo") os.makedirs(self.repodir) - self.sourcedir = os.path.join(self.repodir, "source_protos") - os.makedirs(self.sourcedir) logdir = os.path.join(self.directory, "logs") if casd else None @@ -181,16 +180,29 @@ class ArtifactShare(BaseArtifactShare): channel.close() def get_source_proto(self, source_name): - source_proto = source_pb2.Source() - source_path = os.path.join(self.sourcedir, source_name) - + url = urlparse(self.repo) + channel = grpc.insecure_channel("{}:{}".format(url.hostname, url.port)) try: - with open(source_path, "rb") as f: - source_proto.ParseFromString(f.read()) - except FileNotFoundError: - return None + fetch_service = remote_asset_pb2_grpc.FetchStub(channel) + + uri = REMOTE_ASSET_SOURCE_URN_TEMPLATE.format(source_name) + + request = remote_asset_pb2.FetchDirectoryRequest() + request.uris.append(uri) + + try: + response = fetch_service.FetchDirectory(request) + except grpc.RpcError as e: + if e.code() == grpc.StatusCode.NOT_FOUND: + return None + raise - return source_proto + if response.status.code != code_pb2.OK: + return None + + return response.root_directory_digest + finally: + channel.close() def get_cas_files(self, artifact_proto_digest): |