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")