test wait not called

Author: chuyang-deng

tests/unit/sagemaker/automl/test_auto_ml.py

@@ -294,7 +294,18 @@ def test_auto_ml_only_one_of_problem_type_and_job_objective_provided(sagemaker_s
         )
 
 
-def test_auto_ml_additional_optional_params(sagemaker_session, caplog):
+@patch("sagemaker.automl.automl.AutoMLJob.start_new")
+def test_auto_ml_fit_set_logs_to_false(start_new, sagemaker_session, caplog):
+    auto_ml = AutoML(
+        role=ROLE, target_attribute_name=TARGET_ATTRIBUTE_NAME, sagemaker_session=sagemaker_session
+    )
+    inputs = DEFAULT_S3_INPUT_DATA
+    auto_ml.fit(inputs, job_name=JOB_NAME, wait=False, logs=True)
+    start_new.wait.assert_not_called()
+    assert "Setting logs to False. logs is only meaningful when wait is True." in caplog.text
+
+
+def test_auto_ml_additional_optional_params(sagemaker_session):
     auto_ml = AutoML(
         role=ROLE,
         target_attribute_name=TARGET_ATTRIBUTE_NAME,
@@ -314,7 +325,7 @@ def test_auto_ml_additional_optional_params(sagemaker_session, caplog):
         tags=TAGS,
     )
     inputs = DEFAULT_S3_INPUT_DATA
-    auto_ml.fit(inputs, job_name=JOB_NAME, wait=False, logs=True)
+    auto_ml.fit(inputs, job_name=JOB_NAME)
     sagemaker_session.auto_ml.assert_called_once()
     _, args = sagemaker_session.auto_ml.call_args
 
@@ -348,7 +359,6 @@ def test_auto_ml_additional_optional_params(sagemaker_session, caplog):
         "generate_candidate_definitions_only": GENERATE_CANDIDATE_DEFINITIONS_ONLY,
         "tags": TAGS,
     }
-    assert "logs will be set to False. logs is only meaningful when wait is True." in caplog.text
 
 
 @patch("time.strftime", return_value=TIMESTAMP)