def _stage_files(self, files):
"""Utility method to stage files.
Args:
files: a list of tuples of the form [(local_name, remote_name),...]
describing the name of the artifacts in local temp folder and desired
name in staging location.
"""
server = grpc.server(thread_pool_executor.shared_unbounded_instance())
staging_service = TestLocalFileSystemLegacyArtifactStagingServiceServicer(
self._remote_dir)
beam_artifact_api_pb2_grpc.add_LegacyArtifactStagingServiceServicer_to_server(
staging_service, server)
test_port = server.add_insecure_port('[::]:0')
server.start()
stager = portable_stager.PortableStager(
artifact_service_channel=grpc.insecure_channel('localhost:%s' %
test_port),
staging_session_token='token')
for from_file, to_file in files:
stager.stage_artifact(local_path_to_artifact=os.path.join(
self._temp_dir, from_file),
artifact_name=to_file)
stager.commit_manifest()
return staging_service.manifest.artifact, staging_service.retrieval_tokens
def _start_artifact_service(self, jar, requested_port):
self._artifact_staging_service = artifact_service.ZipFileArtifactService(
jar, self.ARTIFACT_FOLDER)
self._artifact_staging_server = grpc.server(futures.ThreadPoolExecutor())
port = self._artifact_staging_server.add_insecure_port(
'[::]:%s' % requested_port)
beam_artifact_api_pb2_grpc.add_LegacyArtifactStagingServiceServicer_to_server(
self._artifact_staging_service, self._artifact_staging_server)
self._artifact_staging_endpoint = endpoints_pb2.ApiServiceDescriptor(
url='localhost:%d' % port)
self._artifact_staging_server.start()
_LOGGER.info('Artifact server started on port %s', port)
return port
def start_grpc_server(self, port=0):
self._server = grpc.server(UnboundedThreadPoolExecutor())
port = self._server.add_insecure_port(
'%s:%d' % (self.get_bind_address(), port))
beam_job_api_pb2_grpc.add_JobServiceServicer_to_server(self, self._server)
beam_artifact_api_pb2_grpc.add_LegacyArtifactStagingServiceServicer_to_server(
self._artifact_service, self._server)
hostname = self.get_service_address()
self._artifact_staging_endpoint = endpoints_pb2.ApiServiceDescriptor(
url='%s:%d' % (hostname, port))
self._server.start()
_LOGGER.info('Grpc server started at %s on port %d' % (hostname, port))
return port
def test_with_grpc(self):
server = grpc.server(UnboundedThreadPoolExecutor())
try:
beam_artifact_api_pb2_grpc.add_LegacyArtifactStagingServiceServicer_to_server(
self._service, server)
beam_artifact_api_pb2_grpc.add_LegacyArtifactRetrievalServiceServicer_to_server(
self._service, server)
port = server.add_insecure_port('[::]:0')
server.start()
channel = grpc.insecure_channel('localhost:%d' % port)
self._run_staging(
beam_artifact_api_pb2_grpc.LegacyArtifactStagingServiceStub(
channel),
beam_artifact_api_pb2_grpc.LegacyArtifactRetrievalServiceStub(
channel))
channel.close()
finally:
server.stop(1)