Fix formatting

Author: samruds

tests/unit/sagemaker/serve/builder/test_transformers_builder.py

@@ -115,8 +115,7 @@ def test_image_uri(
             model=mock_model_id,
             schema_builder=mock_schema_builder,
             mode=Mode.LOCAL_CONTAINER,
-            vpc_config=MOCK_VPC_CONFIG,
-            image_config=MOCK_IMAGE_CONFIG,
+            image_uri=MOCK_IMAGE_CONFIG,
         )
 
         builder._prepare_for_mode = MagicMock()
@@ -128,17 +127,11 @@ def test_image_uri(
         builder.modes[str(Mode.LOCAL_CONTAINER)] = MagicMock()
         predictor = model.deploy(model_data_download_timeout=1800)
 
-        assert model.image_config == MOCK_IMAGE_CONFIG
-        assert model.vpc_config == MOCK_VPC_CONFIG
+        assert builder.image_uri == MOCK_IMAGE_CONFIG
         assert builder.env_vars["MODEL_LOADING_TIMEOUT"] == "1800"
         assert isinstance(predictor, TransformersLocalModePredictor)
 
         assert builder.nb_instance_type == "ml.g5.24xlarge"
 
-        builder._original_deploy = MagicMock()
-        builder._prepare_for_mode.return_value = (None, {})
-        predictor = model.deploy(mode=Mode.SAGEMAKER_ENDPOINT, role="mock_role_arn")
-        assert "HF_MODEL_ID" in model.env
-
         with self.assertRaises(ValueError) as _:
             model.deploy(mode=Mode.IN_PROCESS)