fix codestyle

Author: mollyheamazon

tests/integ/sagemaker/experiments/test_run.py

@@ -32,7 +32,7 @@
 from sagemaker.experiments.trial_component import _TrialComponent
 from sagemaker.sklearn import SKLearn
 from sagemaker.utils import retry_with_backoff, unique_name_from_base
-from tests.integ.sagemaker.experiments.helpers import name, cleanup_exp_resources, clear_run_context
+from tests.integ.sagemaker.experiments.helpers import name, cleanup_exp_resources
 from sagemaker.experiments.run import (
     RUN_NAME_BASE,
     DELIMITER,
@@ -55,7 +55,7 @@ def artifact_file_path(tempdir):
 metric_name = "Test-Local-Init-Log-Metric"
 
 
-def test_local_run_with_load(sagemaker_session, artifact_file_path, clear_run_context):
+def test_local_run_with_load(sagemaker_session, artifact_file_path):
     exp_name = f"My-Local-Exp-{name()}"
     with cleanup_exp_resources(exp_names=[exp_name], sagemaker_session=sagemaker_session):
         # Run name is not provided, will create a new TC
@@ -87,7 +87,7 @@ def verify_load_run():
 
 
 def test_two_local_run_init_with_same_run_name_and_different_exp_names(
-    sagemaker_session, clear_run_context
+    sagemaker_session
 ):
     exp_name1 = f"my-two-local-exp1-{name()}"
     exp_name2 = f"my-two-local-exp2-{name()}"
@@ -127,7 +127,7 @@ def test_two_local_run_init_with_same_run_name_and_different_exp_names(
     ],
 )
 def test_run_name_vs_trial_component_name_edge_cases(
-    sagemaker_session, input_names, clear_run_context
+    sagemaker_session, input_names
 ):
     exp_name, run_name, run_display_name = input_names
     with cleanup_exp_resources(exp_names=[exp_name], sagemaker_session=sagemaker_session):
@@ -181,7 +181,6 @@ def test_run_from_local_and_train_job_and_all_exp_cfg_match(
     execution_role,
     sagemaker_client_config,
     sagemaker_metrics_config,
-    clear_run_context,
 ):
     # Notes:
     # 1. The 1st Run created locally and its exp config was auto passed to the job
@@ -282,7 +281,6 @@ def test_run_from_local_and_train_job_and_exp_cfg_not_match(
     execution_role,
     sagemaker_client_config,
     sagemaker_metrics_config,
-    clear_run_context,
 ):
     # Notes:
     # 1. The 1st Run created locally and its exp config was auto passed to the job
@@ -369,7 +367,6 @@ def test_run_from_train_job_only(
     execution_role,
     sagemaker_client_config,
     sagemaker_metrics_config,
-    clear_run_context,
 ):
     # Notes:
     # 1. No Run created locally or specified in experiment config
@@ -420,7 +417,6 @@ def test_run_from_processing_job_and_override_default_exp_config(
     execution_role,
     sagemaker_client_config,
     sagemaker_metrics_config,
-    clear_run_context,
 ):
     # Notes:
     # 1. The 1st Run (run) created locally
@@ -500,7 +496,6 @@ def test_run_from_transform_job(
     execution_role,
     sagemaker_client_config,
     sagemaker_metrics_config,
-    clear_run_context,
 ):
     # Notes:
     # 1. The 1st Run (run) created locally
@@ -582,7 +577,6 @@ def test_load_run_auto_pass_in_exp_config_to_job(
     execution_role,
     sagemaker_client_config,
     sagemaker_metrics_config,
-    clear_run_context,
 ):
     # Notes:
     # 1. In local side, load the Run created previously and invoke a job under the load context
@@ -631,7 +625,7 @@ def test_load_run_auto_pass_in_exp_config_to_job(
         )
 
 
-def test_list(run_obj, sagemaker_session, clear_run_context):
+def test_list(run_obj, sagemaker_session):
     tc1 = _TrialComponent.create(
         trial_component_name=f"non-run-tc1-{name()}",
         sagemaker_session=sagemaker_session,
@@ -653,7 +647,7 @@ def test_list(run_obj, sagemaker_session, clear_run_context):
     assert run_tcs[0].experiment_config == run_obj.experiment_config
 
 
-def test_list_twice(run_obj, sagemaker_session, clear_run_context):
+def test_list_twice(run_obj, sagemaker_session):
     tc1 = _TrialComponent.create(
         trial_component_name=f"non-run-tc1-{name()}",
         sagemaker_session=sagemaker_session,

tests/integ/test_feature_store.py

@@ -1645,9 +1645,11 @@ def test_create_dataset_with_feature_group_base(
             feature_store_session, feature_group, offline_store_s3_uri
         )
 
-        with timeout(minutes=10) and cleanup_offline_store(
-            base, feature_store_session
-        ) and cleanup_offline_store(feature_group, feature_store_session):
+        with (
+            timeout(minutes=10) 
+            and cleanup_offline_store(base, feature_store_session) 
+            and cleanup_offline_store(feature_group, feature_store_session)
+        ):
             feature_store = FeatureStore(sagemaker_session=feature_store_session)
             df, query_string = (
                 feature_store.create_dataset(base=base, output_path=offline_store_s3_uri)
@@ -1832,9 +1834,11 @@ def test_create_dataset_with_feature_group_base_with_additional_params(
             feature_store_session, feature_group, offline_store_s3_uri
         )
 
-        with timeout(minutes=10) and cleanup_offline_store(
-            base, feature_store_session
-        ) and cleanup_offline_store(feature_group, feature_store_session):
+        with (
+            timeout(minutes=10) 
+            and cleanup_offline_store(base, feature_store_session) 
+            and cleanup_offline_store(feature_group, feature_store_session)
+        ):
             feature_store = FeatureStore(sagemaker_session=feature_store_session)
             df, query_string = (
                 feature_store.create_dataset(base=base, output_path=offline_store_s3_uri)

tests/unit/sagemaker/modules/train/test_model_trainer.py

@@ -1050,17 +1050,17 @@ def mock_upload_data(path, bucket, key_prefix):
     model_trainer.train()
 
     mock_local_container.assert_called_once_with(
-      training_job_name=unique_name,
-      instance_type=compute.instance_type,
-      instance_count=compute.instance_count,
-      image=training_image,
-      container_root=local_container_root,
-      sagemaker_session=modules_session,
-      container_entrypoint=DEFAULT_ENTRYPOINT,
-      container_arguments=DEFAULT_ARGUMENTS,
-      input_data_config=ANY,
-      hyper_parameters=hyperparameters,
-      environment=environment,
+        training_job_name=unique_name,
+        instance_type=compute.instance_type,
+        instance_count=compute.instance_count,
+        image=training_image,
+        container_root=local_container_root,
+        sagemaker_session=modules_session,
+        container_entrypoint=DEFAULT_ENTRYPOINT,
+        container_arguments=DEFAULT_ARGUMENTS,
+        input_data_config=ANY,
+        hyper_parameters=hyperparameters,
+        environment=environment,
 )
 
 

tests/unit/sagemaker/workflow/test_pipeline.py

@@ -19,7 +19,7 @@
 import pytest
 
 from mock import Mock, call, patch
-from mock.mock import MagicMock, ANY
+from mock.mock import MagicMock
 
 from sagemaker import s3
 from sagemaker.remote_function.job import _JobSettings
@@ -345,7 +345,7 @@ def test_pipeline_update_with_parallelism_config(sagemaker_session_mock, role_ar
         role_arn=role_arn,
         parallelism_config=dict(MaxParallelExecutionSteps=10),
     )
-    pipeline.update(  
+    pipeline.update(
         role_arn=role_arn,
         parallelism_config={"MaxParallelExecutionSteps": 10},
     )
@@ -422,9 +422,7 @@ def _raise_does_already_exists_client_error(**kwargs):
     sagemaker_session_mock.sagemaker_client.update_pipeline.assert_called_once_with(
         PipelineName="MyPipeline", PipelineDefinition=pipeline.definition(), RoleArn=role_arn
     )
-    sagemaker_session_mock.sagemaker_client.list_tags.assert_called_with(
-        ResourceArn="pipeline-arn"
-    )
+    sagemaker_session_mock.sagemaker_client.list_tags.assert_called_with(ResourceArn="pipeline-arn")
 
     tags.append({"Key": "dummy", "Value": "dummy_tag"})
     sagemaker_session_mock.sagemaker_client.add_tags.assert_called_with(