def _get_model_version_stages():
request_message = _get_request_message(GetModelVersionStages())
stages = _get_model_registry_store().get_model_version_stages(
ModelVersion.from_proto(request_message.model_version))
response_message = GetModelVersionStages.Response()
response_message.stages.extend(stages)
return _wrap_response(response_message)
def get_model_version_stages(self, model_version):
"""
:return: A list of valid stages.
"""
req_body = message_to_json(GetModelVersionStages(model_version=model_version.to_proto()))
response_proto = self._call_endpoint(GetModelVersionStages, req_body)
return response_proto.stages
def test_get_model_version_stages(self, mock_http):
rm = RegisteredModel("model_11")
mv = ModelVersion(rm, 8)
self.store.get_model_version_stages(model_version=mv)
self._verify_requests(
mock_http, "model-versions/get-stages", "POST",
GetModelVersionStages(model_version=mv.to_proto()))
def test_model_version_stages(mock_get_request_message, mock_model_registry_store):
rm = RegisteredModel("model1")
mv = ModelVersion(registered_model=rm, version=32)
mock_get_request_message.return_value = GetModelVersionStages(model_version=mv.to_proto())
stages = ["Stage1", "Production", "0", "5% traffic", "None"]
mock_model_registry_store.get_model_version_stages.return_value = stages
resp = _get_model_version_stages()
args, _ = mock_model_registry_store.get_model_version_stages.call_args
assert args == (mv, )
assert json.loads(resp.get_data()) == {"stages": stages}