Fix unit tests

aws · Dec 5, 2024 · ec5f37d · ec5f37d
1 parent 213b969
commit ec5f37d
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tests/unit/sagemaker/modules/train/test_model_trainer.py b/tests/unit/sagemaker/modules/train/test_model_trainer.py
@@ -285,6 +285,7 @@ def test_train_with_intelligent_defaults_training_job_space(
             volume_kms_key_id=None,
             keep_alive_period_in_seconds=None,
             instance_groups=None,
+            training_plan_arn=None,
         ),
         vpc_config=None,
         session=ANY,
@@ -825,6 +826,7 @@ def mock_upload_data(path, bucket, key_prefix):
             volume_kms_key_id=compute.volume_kms_key_id,
             keep_alive_period_in_seconds=compute.keep_alive_period_in_seconds,
             instance_groups=None,
+            training_plan_arn=None,
         ),
         vpc_config=VpcConfig(
             security_group_ids=networking.security_group_ids,