Пример #1
0
def test_predictor_classify(sagemaker_session):
    predictor = TensorFlowPredictor("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
Пример #2
0
def test_predictor_classify_bad_content_type(sagemaker_session):
    predictor = TensorFlowPredictor("endpoint", sagemaker_session,
                                    CSVSerializer())

    with pytest.raises(ValueError):
        predictor.classify(CLASSIFY_INPUT)