def test_predictor_regress(sagemaker_session):
predictor = TensorFlowPredictor(
"endpoint", sagemaker_session, model_name="model", model_version="123"
)
mock_response(json.dumps(REGRESS_RESPONSE).encode("utf-8"), sagemaker_session)
result = predictor.regress(REGRESS_INPUT)
assert_invoked_with_body_dict(
sagemaker_session,
EndpointName="endpoint",
ContentType=JSON_CONTENT_TYPE,
Accept=JSON_CONTENT_TYPE,
CustomAttributes="tfs-method=regress,tfs-model-name=model,tfs-model-version=123",
Body=json.dumps(REGRESS_INPUT),
)
assert REGRESS_RESPONSE == result
def test_predictor_regress_bad_content_type(sagemaker_session):
predictor = TensorFlowPredictor("endpoint", sagemaker_session,
CSVSerializer())
with pytest.raises(ValueError):
predictor.regress(REGRESS_INPUT)