def test_predictor_classify(sagemaker_session):
predictor = Predictor('endpoint', sagemaker_session)
mock_response(json.dumps(CLASSIFY_RESPONSE).encode('utf-8'), sagemaker_session)
result = predictor.classify(CLASSIFY_INPUT)
assert_invoked_with_body_dict(sagemaker_session,
EndpointName='endpoint',
ContentType=JSON_CONTENT_TYPE,
Accept=JSON_CONTENT_TYPE,
CustomAttributes='tfs-method=classify',
Body=json.dumps(CLASSIFY_INPUT))
assert CLASSIFY_RESPONSE == result
def test_predictor_classify_bad_content_type(sagemaker_session):
predictor = Predictor("endpoint", sagemaker_session, csv_serializer)
with pytest.raises(ValueError):
predictor.classify(CLASSIFY_INPUT)
