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)