black check

Author: imujjwal96

tests/integ/test_tf_script_mode.py

@@ -159,7 +159,7 @@ def test_mnist_async(sagemaker_session):
     training_job_name = estimator.latest_training_job.name
     time.sleep(20)
     endpoint_name = training_job_name
-    model_name = 'model-name-1'
+    model_name = "model-name-1"
     _assert_training_job_tags_match(
         sagemaker_session.sagemaker_client, estimator.latest_training_job.name, TAGS
     )
@@ -168,8 +168,10 @@ def test_mnist_async(sagemaker_session):
             training_job_name=training_job_name, sagemaker_session=sagemaker_session
         )
         predictor = estimator.deploy(
-            initial_instance_count=1, instance_type="ml.c4.xlarge", endpoint_name=endpoint_name,
-            model_name=model_name
+            initial_instance_count=1,
+            instance_type="ml.c4.xlarge",
+            endpoint_name=endpoint_name,
+            model_name=model_name,
         )
 
         result = predictor.predict(np.zeros(784))
@@ -178,9 +180,7 @@ def test_mnist_async(sagemaker_session):
         _assert_model_tags_match(
             sagemaker_session.sagemaker_client, estimator.latest_training_job.name, TAGS
         )
-        _assert_model_name_match(
-            sagemaker_session.sagemaker_client, endpoint_name, model_name
-        )
+        _assert_model_name_match(sagemaker_session.sagemaker_client, endpoint_name, model_name)
 
 
 def test_deploy_with_input_handlers(sagemaker_session, instance_type):
@@ -252,4 +252,4 @@ def _assert_model_name_match(sagemaker_client, endpoint_config_name, model_name)
     endpoint_config_description = sagemaker_client.describe_endpoint_config(
         EndpointConfigName=endpoint_config_name
     )
-    assert model_name == endpoint_config_description['ProductionVariants'][0]['ModelName']
+    assert model_name == endpoint_config_description["ProductionVariants"][0]["ModelName"]

tests/integ/test_tuner.py

@@ -844,7 +844,7 @@ def test_attach_tuning_pytorch(sagemaker_session):
         tuner.wait()
 
     endpoint_name = tuning_job_name
-    model_name = 'model-name-1'
+    model_name = "model-name-1"
     attached_tuner = HyperparameterTuner.attach(
         tuning_job_name, sagemaker_session=sagemaker_session
     )
@@ -948,5 +948,7 @@ def _fm_serializer(data):
 
 
 def _assert_model_name_match(sagemaker_client, endpoint_config_name, model_name):
-    endpoint_config_description = sagemaker_client.describe_endpoint_config(EndpointConfigName=endpoint_config_name)
-    assert model_name == endpoint_config_description['ProductionVariants'][0]['ModelName']
+    endpoint_config_description = sagemaker_client.describe_endpoint_config(
+        EndpointConfigName=endpoint_config_name
+    )
+    assert model_name == endpoint_config_description["ProductionVariants"][0]["ModelName"]

tests/unit/test_estimator.py

@@ -1712,8 +1712,14 @@ def test_deploy_with_update_endpoint(sagemaker_session):
 
 
 def test_deploy_with_model_name(sagemaker_session):
-    estimator = Estimator(IMAGE_NAME, ROLE, INSTANCE_COUNT, INSTANCE_TYPE, output_path=OUTPUT_PATH,
-                          sagemaker_session=sagemaker_session)
+    estimator = Estimator(
+        IMAGE_NAME,
+        ROLE,
+        INSTANCE_COUNT,
+        INSTANCE_TYPE,
+        output_path=OUTPUT_PATH,
+        sagemaker_session=sagemaker_session,
+    )
     estimator.set_hyperparameters(**HYPERPARAMS)
     estimator.fit({"train": "s3://bucket/training-prefix"})
     model_name = "model-name"
@@ -1725,10 +1731,16 @@ def test_deploy_with_model_name(sagemaker_session):
 
 
 def test_deploy_with_no_model_name(sagemaker_session):
-    estimator = Estimator(IMAGE_NAME, ROLE, INSTANCE_COUNT, INSTANCE_TYPE, output_path=OUTPUT_PATH,
-                          sagemaker_session=sagemaker_session)
+    estimator = Estimator(
+        IMAGE_NAME,
+        ROLE,
+        INSTANCE_COUNT,
+        INSTANCE_TYPE,
+        output_path=OUTPUT_PATH,
+        sagemaker_session=sagemaker_session,
+    )
     estimator.set_hyperparameters(**HYPERPARAMS)
-    estimator.fit({'train': 's3://bucket/training-prefix'})
+    estimator.fit({"train": "s3://bucket/training-prefix"})
     estimator.deploy(INSTANCE_COUNT, INSTANCE_TYPE)
 
     sagemaker_session.create_model.assert_called_once()

tests/unit/test_tuner.py

@@ -646,7 +646,7 @@ def test_deploy_default(tuner):
     tuner.estimator.sagemaker_session.create_model.assert_called_once()
     args = tuner.estimator.sagemaker_session.create_model.call_args[0]
 
-    assert args[0] == 'neo'
+    assert args[0] == "neo"
     assert args[1] == ROLE
     assert args[2]["Image"] == IMAGE_NAME
     assert args[2]["ModelDataUrl"] == MODEL_DATA