Improvements of the integration tests

Author: Anton Repushko

src/sagemaker/automl/automlv2.py

@@ -877,8 +877,8 @@ def best_candidate(self, job_name=None):
         """Returns the best candidate of an AutoML job for a given name.
 
         Args:
-            job_name (str): The name of the AutoML job. If None, will use object's
-                _current_auto_ml_job_name.
+            job_name (str): The name of the AutoML job. If None, object's
+                _current_auto_ml_job_name will be used.
 
         Returns:
             dict: A dictionary with information of the best candidate.

tests/integ/auto_ml_v2_utils.py

@@ -69,6 +69,7 @@
         item_identifier_attribute_name="item_id",
         target_attribute_name="target",
         timestamp_attribute_name="timestamp",
+        forecast_quantiles=["p10", "p50", "p90"],
     ),
 }
 

tests/integ/test_auto_ml_v2.py

@@ -174,7 +174,9 @@ def test_auto_ml_v2_attach(problem_type, job_name_fixture_key, sagemaker_session
     assert desc["AutoMLJobName"] == job_name
     assert desc["AutoMLJobStatus"] in ["InProgress", "Completed"]
     assert desc["AutoMLJobSecondaryStatus"] != "Failed"
-    assert desc["ProblemConfig"] == auto_ml_v2_utils.PROBLEM_CONFIGS[problem_type]
+    assert sorted(desc["AutoMLProblemTypeConfig"]) == sorted(
+        auto_ml_v2_utils.PROBLEM_CONFIGS[problem_type]
+    )
     assert desc["OutputDataConfig"] == expected_default_output_config
 
 
@@ -251,6 +253,8 @@ def test_list_candidates(
 
         candidates = auto_ml.list_candidates(job_name=job_name)
         assert len(candidates) == num_candidates
+    else:
+        pytest.skip("The job hasn't finished yet")
 
 
 @pytest.mark.skipif(
@@ -320,6 +324,8 @@ def test_best_candidate(
         best_candidate = auto_ml.best_candidate(job_name=job_name)
         assert len(best_candidate["InferenceContainers"]) == num_containers
         assert best_candidate["CandidateStatus"] == "Completed"
+    else:
+        pytest.skip("The job hasn't finished yet")
 
 
 @pytest.mark.skipif(
@@ -411,6 +417,8 @@ def test_deploy_best_candidate(
         )["EndpointStatus"]
         assert endpoint_status == "InService"
         sagemaker_session.sagemaker_client.delete_endpoint(EndpointName=endpoint_name)
+    else:
+        pytest.skip("The job hasn't finished yet")
 
 
 @pytest.mark.skipif(
@@ -482,3 +490,5 @@ def test_candidate_estimator_get_steps(
         candidate_estimator = CandidateEstimator(candidate, sagemaker_session)
         steps = candidate_estimator.get_steps()
         assert len(steps) == num_steps
+    else:
+        pytest.skip("The job hasn't finished yet")