fix: fix broken unit tests due to refactoring

Author: Balaji Sankar

src/sagemaker/local/local_session.py

@@ -682,8 +682,9 @@ def _initialize(
             self.s3_resource = boto_session.resource("s3", endpoint_url=self.s3_endpoint_url)
             self.s3_client = boto_session.client("s3", endpoint_url=self.s3_endpoint_url)
             self.sagemaker_config = (
-                sagemaker_config if sagemaker_config else fetch_sagemaker_config(
-                    s3_resource=self.s3_resource)
+                sagemaker_config
+                if sagemaker_config
+                else fetch_sagemaker_config(s3_resource=self.s3_resource)
             )
         else:
             self.sagemaker_config = (

tests/conftest.py

@@ -166,7 +166,7 @@ def sagemaker_session(
         sagemaker_client=sagemaker_client,
         sagemaker_runtime_client=runtime_client,
         sagemaker_metrics_client=metrics_client,
-        sagemaker_config={"SchemaVersion": "1.0"},
+        sagemaker_config={},
     )
 
 

tests/unit/sagemaker/model/test_deploy.py

@@ -168,6 +168,7 @@ def test_deploy_accelerator_type(
 @patch("sagemaker.model.Model._create_sagemaker_model", Mock())
 @patch("sagemaker.production_variant", return_value=BASE_PRODUCTION_VARIANT)
 def test_deploy_endpoint_name(sagemaker_session):
+    sagemaker_session.sagemaker_config = {}
     model = Model(MODEL_IMAGE, MODEL_DATA, role=ROLE, sagemaker_session=sagemaker_session)
 
     endpoint_name = "blah"
@@ -371,6 +372,7 @@ def test_deploy_async_inference(production_variant, name_from_base, sagemaker_se
 @patch("sagemaker.model.Model._create_sagemaker_model")
 @patch("sagemaker.production_variant")
 def test_deploy_serverless_inference(production_variant, create_sagemaker_model, sagemaker_session):
+    sagemaker_session.sagemaker_config = {}
     model = Model(
         MODEL_IMAGE, MODEL_DATA, role=ROLE, name=MODEL_NAME, sagemaker_session=sagemaker_session
     )
@@ -439,8 +441,8 @@ def test_deploy_wrong_serverless_config(sagemaker_session):
 @patch("sagemaker.session.Session")
 @patch("sagemaker.local.LocalSession")
 def test_deploy_creates_correct_session(local_session, session):
-    local_session.sagemaker_config = {}
-    session.sagemaker_config = {}
+    local_session.return_value.sagemaker_config = {}
+    session.return_value.sagemaker_config = {}
     # We expect a LocalSession when deploying to instance_type = 'local'
     model = Model(MODEL_IMAGE, MODEL_DATA, role=ROLE)
     model.deploy(endpoint_name="blah", instance_type="local", initial_instance_count=1)

tests/unit/sagemaker/model/test_framework_model.py

@@ -208,6 +208,7 @@ def test_git_support_repo_not_provided(sagemaker_session):
     ),
 )
 def test_git_support_git_clone_fail(sagemaker_session):
+    sagemaker_session.sagemaker_config = {}
     entry_point = "source_dir/entry_point"
     git_config = {"repo": "https://github.com/aws/no-such-repo.git", "branch": BRANCH}
     with pytest.raises(subprocess.CalledProcessError) as error:
@@ -257,6 +258,7 @@ def test_git_support_commit_not_exist(git_clone_repo, sagemaker_session):
     side_effect=ValueError("Entry point does not exist in the repo."),
 )
 def test_git_support_entry_point_not_exist(sagemaker_session):
+    sagemaker_session.sagemaker_config = {}
     entry_point = "source_dir/entry_point"
     git_config = {"repo": GIT_REPO, "branch": BRANCH, "commit": COMMIT}
     with pytest.raises(ValueError) as error:
@@ -272,6 +274,7 @@ def test_git_support_entry_point_not_exist(sagemaker_session):
     side_effect=ValueError("Source directory does not exist in the repo."),
 )
 def test_git_support_source_dir_not_exist(sagemaker_session):
+    sagemaker_session.sagemaker_config = {}
     entry_point = "entry_point"
     source_dir = "source_dir_that_does_not_exist"
     git_config = {"repo": GIT_REPO, "branch": BRANCH, "commit": COMMIT}
@@ -291,6 +294,7 @@ def test_git_support_source_dir_not_exist(sagemaker_session):
     side_effect=ValueError("Dependency no-such-dir does not exist in the repo."),
 )
 def test_git_support_dependencies_not_exist(sagemaker_session):
+    sagemaker_session.sagemaker_config = {}
     entry_point = "entry_point"
     dependencies = ["foo", "no_such_dir"]
     git_config = {"repo": GIT_REPO, "branch": BRANCH, "commit": COMMIT}

tests/unit/sagemaker/model/test_model.py

@@ -326,6 +326,8 @@ def test_create_sagemaker_model_generates_model_name_each_time(
 @patch("sagemaker.session.Session")
 @patch("sagemaker.local.LocalSession")
 def test_create_sagemaker_model_creates_correct_session(local_session, session):
+    local_session.return_value.sagemaker_config = {}
+    session.return_value.sagemaker_config = {}
     model = Model(MODEL_IMAGE, MODEL_DATA)
     model._create_sagemaker_model("local")
     assert model.sagemaker_session == local_session.return_value
@@ -433,6 +435,8 @@ def test_model_create_transformer_base_name(sagemaker_session):
 @patch("sagemaker.session.Session")
 @patch("sagemaker.local.LocalSession")
 def test_transformer_creates_correct_session(local_session, session):
+    local_session.return_value.sagemaker_config = {}
+    session.return_value.sagemaker_config = {}
     model = Model(MODEL_IMAGE, MODEL_DATA, sagemaker_session=None)
     transformer = model.transformer(instance_count=1, instance_type="local")
     assert model.sagemaker_session == local_session.return_value

tests/unit/sagemaker/model/test_neo.py

@@ -170,6 +170,7 @@ def test_compile_model_for_cloud_tflite(sagemaker_session):
 @patch("sagemaker.session.Session")
 def test_compile_creates_session(session):
     session.return_value.boto_region_name = REGION
+    session.return_value.sagemaker_config = {}
 
     model = _create_model()
     model.compile(
@@ -313,6 +314,7 @@ def test_compile_with_framework_version_16(sagemaker_session):
 @patch("sagemaker.session.Session")
 def test_compile_with_pytorch_neo_in_ml_inf(session):
     session.return_value.boto_region_name = REGION
+    session.return_value.sagemaker_config = {}
 
     model = _create_model()
     model.compile(
@@ -336,6 +338,7 @@ def test_compile_with_pytorch_neo_in_ml_inf(session):
 @patch("sagemaker.session.Session")
 def test_compile_with_tensorflow_neo_in_ml_inf(session):
     session.return_value.boto_region_name = REGION
+    session.return_value.sagemaker_config = {}
 
     model = _create_model()
     model.compile(

tests/unit/sagemaker/monitor/test_model_monitoring.py

@@ -877,8 +877,14 @@ def test_data_quality_batch_transform_monitor_create_schedule_with_sagemaker_con
     data_quality_monitor,
     sagemaker_session,
 ):
+    from sagemaker.utils import get_config_value
 
     sagemaker_session.sagemaker_config = SAGEMAKER_CONFIG_MONITORING_SCHEDULE
+    sagemaker_session._append_sagemaker_config_tags = Mock(
+        name="_append_sagemaker_config_tags",
+        side_effect=lambda tags, config_path_to_tags: tags
+        + get_config_value(config_path_to_tags, SAGEMAKER_CONFIG_MONITORING_SCHEDULE),
+    )
 
     sagemaker_session.sagemaker_client.create_monitoring_schedule = Mock()
     data_quality_monitor.sagemaker_session = sagemaker_session

tests/unit/sagemaker/spark/test_processing.py

@@ -61,6 +61,7 @@ def sagemaker_session():
         settings=SessionSettings(),
     )
     session_mock.default_bucket = Mock(name="default_bucket", return_value=BUCKET_NAME)
+    session_mock.sagemaker_config = {}
 
     return session_mock
 

tests/unit/sagemaker/workflow/test_pipeline.py

@@ -71,8 +71,10 @@ def test_pipeline_create_and_update_without_role_arn(sagemaker_session_mock):
 
 def test_pipeline_create_and_update_with_config_injection(sagemaker_session_mock):
     # For tests which doesn't verify config file injection, operate with empty config
+    pipeline_role_arn = "arn:aws:iam::111111111111:role/ConfigRole"
     sagemaker_session_mock.sagemaker_config = {
-        "SageMaker": {"Pipeline": {"RoleArn": "ConfigRoleArn"}}
+        "SchemaVersion": "1.0",
+        "SageMaker": {"Pipeline": {"RoleArn": pipeline_role_arn}},
     }
     sagemaker_session_mock.sagemaker_client.describe_pipeline.return_value = {
         "PipelineArn": "pipeline-arn"
@@ -85,15 +87,21 @@ def test_pipeline_create_and_update_with_config_injection(sagemaker_session_mock
     )
     pipeline.create()
     sagemaker_session_mock.sagemaker_client.create_pipeline.assert_called_with(
-        PipelineName="MyPipeline", PipelineDefinition=pipeline.definition(), RoleArn="ConfigRoleArn"
+        PipelineName="MyPipeline",
+        PipelineDefinition=pipeline.definition(),
+        RoleArn=pipeline_role_arn,
     )
     pipeline.update()
     sagemaker_session_mock.sagemaker_client.update_pipeline.assert_called_with(
-        PipelineName="MyPipeline", PipelineDefinition=pipeline.definition(), RoleArn="ConfigRoleArn"
+        PipelineName="MyPipeline",
+        PipelineDefinition=pipeline.definition(),
+        RoleArn=pipeline_role_arn,
     )
     pipeline.upsert()
     assert sagemaker_session_mock.sagemaker_client.update_pipeline.called_with(
-        PipelineName="MyPipeline", PipelineDefinition=pipeline.definition(), RoleArn="ConfigRoleArn"
+        PipelineName="MyPipeline",
+        PipelineDefinition=pipeline.definition(),
+        RoleArn=pipeline_role_arn,
     )
 
 
@@ -129,6 +137,7 @@ def test_pipeline_create_with_parallelism_config(sagemaker_session_mock, role_ar
 
 @patch("sagemaker.s3.S3Uploader.upload_string_as_file_body")
 def test_large_pipeline_create(sagemaker_session_mock, role_arn):
+    sagemaker_session_mock.sagemaker_config = {}
     parameter = ParameterString("MyStr")
     pipeline = Pipeline(
         name="MyPipeline",
@@ -151,6 +160,7 @@ def test_large_pipeline_create(sagemaker_session_mock, role_arn):
 
 
 def test_pipeline_update(sagemaker_session_mock, role_arn):
+    sagemaker_session_mock.sagemaker_config = {}
     pipeline = Pipeline(
         name="MyPipeline",
         parameters=[],
@@ -201,6 +211,7 @@ def test_pipeline_update_with_parallelism_config(sagemaker_session_mock, role_ar
 
 @patch("sagemaker.s3.S3Uploader.upload_string_as_file_body")
 def test_large_pipeline_update(sagemaker_session_mock, role_arn):
+    sagemaker_session_mock.sagemaker_config = {}
     parameter = ParameterString("MyStr")
     pipeline = Pipeline(
         name="MyPipeline",

tests/unit/sagemaker/workflow/test_utils.py

@@ -226,9 +226,11 @@ def test_inject_repack_script_s3(estimator, tmp, fake_s3):
     model_data = Properties(step_name="MyStep", shape_name="DescribeModelOutput")
     entry_point = "inference.py"
     source_dir_path = "s3://fake/location"
+    session_mock = fake_s3.sagemaker_session
+    session_mock.sagemaker_config = {}
     step = _RepackModelStep(
         name="MyRepackModelStep",
-        sagemaker_session=fake_s3.sagemaker_session,
+        sagemaker_session=session_mock,
         role=estimator.role,
         image_uri="foo",
         model_data=model_data,