def test_validate_and_update_inference_response_wrong_input(): cic = copy.copy(CLASSIFICATION_INFERENCE_CONTAINERS) with pytest.raises( ValueError, message="Requested inference output keys [wrong_key, wrong_label] are unsupported. " "The supported inference keys are [probability, probabilities, predicted_label, labels]", ): AutoML.validate_and_update_inference_response( inference_containers=cic, inference_response_keys=["wrong_key", "wrong_label", "probabilities", "probability"], )
def test_validate_and_update_inference_response(): cic = copy.copy(CLASSIFICATION_INFERENCE_CONTAINERS) AutoML.validate_and_update_inference_response( inference_containers=cic, inference_response_keys=[ "predicted_label", "labels", "probabilities", "probability" ], ) assert (cic[2]["Environment"]["SAGEMAKER_INFERENCE_OUTPUT"] == "predicted_label,labels,probabilities,probability") assert (cic[2]["Environment"]["SAGEMAKER_INFERENCE_INPUT"] == "predicted_label,probabilities,probability") assert (cic[1]["Environment"]["SAGEMAKER_INFERENCE_OUTPUT"] == "predicted_label,probabilities,probability")