@@ -477,16 +477,19 @@ def test_deploy(sagemaker_session, candidate_mock):
477
477
)
478
478
479
479
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
+
481
484
auto_ml = AutoML (
482
485
role = ROLE , target_attribute_name = TARGET_ATTRIBUTE_NAME , sagemaker_session = sagemaker_session
483
486
)
484
- auto_ml .best_candidate = Mock (name = "best_candidate" , return_value = CANDIDATE_DICT )
485
487
auto_ml ._deploy_inference_pipeline = Mock ("_deploy_inference_pipeline" , return_value = None )
486
488
487
489
auto_ml .deploy (
488
490
initial_instance_count = INSTANCE_COUNT ,
489
491
instance_type = INSTANCE_TYPE ,
492
+ candidate = CANDIDATE_DICT ,
490
493
sagemaker_session = sagemaker_session ,
491
494
name = JOB_NAME ,
492
495
endpoint_name = JOB_NAME ,
@@ -515,6 +518,8 @@ def test_deploy_optional_args(sagemaker_session, candidate_mock):
515
518
predictor_cls = RealTimePredictor ,
516
519
)
517
520
521
+ candidate_estimator .assert_called_with (CANDIDATE_DICT , sagemaker_session = sagemaker_session )
522
+
518
523
519
524
def test_candidate_estimator_get_steps (sagemaker_session ):
520
525
candidate_estimator = CandidateEstimator (CANDIDATE_DICT , sagemaker_session = sagemaker_session )
0 commit comments