Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def create_yatai_service_mock(repo_storage_type=BentoUri.LOCAL):
bento_pb = Bento(name='bento_test_name', version='version1.1.1')
if repo_storage_type == BentoUri.LOCAL:
bento_pb.uri.uri = '/tmp/path/to/bundle'
bento_pb.uri.type = repo_storage_type
api = BentoServiceMetadata.BentoServiceApi(name='predict')
bento_pb.bento_service_metadata.apis.extend([api])
bento_pb.bento_service_metadata.env.python_version = '3.7.0'
get_bento_response = GetBentoResponse(bento=bento_pb)
yatai_service_mock = MagicMock()
yatai_service_mock.GetBento.return_value = get_bento_response
return yatai_service_mock
def create_yatai_service_mock(repo_storage_type=BentoUri.LOCAL):
bento_pb = Bento(name='bento_test_name', version='version1.1.1')
if repo_storage_type == BentoUri.LOCAL:
bento_pb.uri.uri = '/fake/path/to/bundle'
bento_pb.uri.type = repo_storage_type
api = BentoServiceMetadata.BentoServiceApi(name='predict')
bento_pb.bento_service_metadata.apis.extend([api])
get_bento_response = GetBentoResponse(bento=bento_pb)
yatai_service_mock = MagicMock()
yatai_service_mock.GetBento.return_value = get_bento_response
return yatai_service_mock
def create_yatai_service_mock(repo_storage_type=BentoUri.LOCAL):
bento_pb = Bento(
name=TEST_DEPLOYMENT_BENTO_NAME, version=TEST_DEPLOYMENT_BENTO_VERSION
)
if repo_storage_type == BentoUri.LOCAL:
bento_pb.uri.uri = TEST_DEPLOYMENT_BENTO_LOCAL_URI
bento_pb.uri.type = repo_storage_type
api = BentoServiceMetadata.BentoServiceApi(name=TEST_BENTO_API_NAME)
bento_pb.bento_service_metadata.apis.extend([api])
get_bento_response = GetBentoResponse(bento=bento_pb)
yatai_service_mock = MagicMock()
yatai_service_mock.GetBento.return_value = get_bento_response
return yatai_service_mock
def GetBento(self, request, context=None):
try:
# TODO: validate request
bento_metadata_pb = self.bento_metadata_store.get(
request.bento_name, request.bento_version
)
if bento_metadata_pb:
return GetBentoResponse(status=Status.OK(), bento=bento_metadata_pb)
else:
return GetBentoResponse(
status=Status.NOT_FOUND(
"Bento `{}:{}` is not found".format(
request.bento_name, request.bento_version
)
)
)
except BentoMLException as e:
logger.error("RPC ERROR GetBento: %s", e)
return GetBentoResponse(status=e.status_proto)