@@ -477,16 +477,19 @@ def test_deploy(sagemaker_session, candidate_mock):
477477 )
478478
479479
480- def test_deploy_optional_args (sagemaker_session , candidate_mock ):
480+ @patch ("sagemaker.automl.automl.CandidateEstimator" )
481+ def test_deploy_optional_args (candidate_estimator , sagemaker_session , candidate_mock ):
482+ candidate_estimator .return_value = candidate_mock
483+
481484 auto_ml = AutoML (
482485 role = ROLE , target_attribute_name = TARGET_ATTRIBUTE_NAME , sagemaker_session = sagemaker_session
483486 )
484- auto_ml .best_candidate = Mock (name = "best_candidate" , return_value = CANDIDATE_DICT )
485487 auto_ml ._deploy_inference_pipeline = Mock ("_deploy_inference_pipeline" , return_value = None )
486488
487489 auto_ml .deploy (
488490 initial_instance_count = INSTANCE_COUNT ,
489491 instance_type = INSTANCE_TYPE ,
492+ candidate = CANDIDATE_DICT ,
490493 sagemaker_session = sagemaker_session ,
491494 name = JOB_NAME ,
492495 endpoint_name = JOB_NAME ,
@@ -515,6 +518,8 @@ def test_deploy_optional_args(sagemaker_session, candidate_mock):
515518 predictor_cls = RealTimePredictor ,
516519 )
517520
521+ candidate_estimator .assert_called_with (CANDIDATE_DICT , sagemaker_session = sagemaker_session )
522+
518523
519524def test_candidate_estimator_get_steps (sagemaker_session ):
520525 candidate_estimator = CandidateEstimator (CANDIDATE_DICT , sagemaker_session = sagemaker_session )
0 commit comments