@@ -652,11 +652,11 @@ def test_s3_input_all_arguments():
652652
653653@pytest .fixture ()
654654def sagemaker_session ():
655- boto_mock = Mock (name = "boto_session" )
655+ boto_mock = MagicMock (name = "boto_session" )
656656 boto_mock .client ("sts" , endpoint_url = STS_ENDPOINT ).get_caller_identity .return_value = {
657657 "Account" : "123"
658658 }
659- ims = sagemaker .Session (boto_session = boto_mock , sagemaker_client = Mock ())
659+ ims = sagemaker .Session (boto_session = boto_mock , sagemaker_client = MagicMock ())
660660 ims .expand_role = Mock (return_value = EXPANDED_ROLE )
661661 return ims
662662
@@ -1352,10 +1352,10 @@ def __init__(self, code):
13521352
13531353@pytest .fixture ()
13541354def sagemaker_session_complete ():
1355- boto_mock = Mock (name = "boto_session" )
1355+ boto_mock = MagicMock (name = "boto_session" )
13561356 boto_mock .client ("logs" ).describe_log_streams .return_value = DEFAULT_LOG_STREAMS
13571357 boto_mock .client ("logs" ).get_log_events .side_effect = DEFAULT_LOG_EVENTS
1358- ims = sagemaker .Session (boto_session = boto_mock , sagemaker_client = Mock ())
1358+ ims = sagemaker .Session (boto_session = boto_mock , sagemaker_client = MagicMock ())
13591359 ims .sagemaker_client .describe_training_job .return_value = COMPLETED_DESCRIBE_JOB_RESULT
13601360 ims .sagemaker_client .describe_transform_job .return_value = (
13611361 COMPLETED_DESCRIBE_TRANSFORM_JOB_RESULT
@@ -1365,21 +1365,21 @@ def sagemaker_session_complete():
13651365
13661366@pytest .fixture ()
13671367def sagemaker_session_stopped ():
1368- boto_mock = Mock (name = "boto_session" )
1368+ boto_mock = MagicMock (name = "boto_session" )
13691369 boto_mock .client ("logs" ).describe_log_streams .return_value = DEFAULT_LOG_STREAMS
13701370 boto_mock .client ("logs" ).get_log_events .side_effect = DEFAULT_LOG_EVENTS
1371- ims = sagemaker .Session (boto_session = boto_mock , sagemaker_client = Mock ())
1371+ ims = sagemaker .Session (boto_session = boto_mock , sagemaker_client = MagicMock ())
13721372 ims .sagemaker_client .describe_training_job .return_value = STOPPED_DESCRIBE_JOB_RESULT
13731373 ims .sagemaker_client .describe_transform_job .return_value = STOPPED_DESCRIBE_TRANSFORM_JOB_RESULT
13741374 return ims
13751375
13761376
13771377@pytest .fixture ()
13781378def sagemaker_session_ready_lifecycle ():
1379- boto_mock = Mock (name = "boto_session" )
1379+ boto_mock = MagicMock (name = "boto_session" )
13801380 boto_mock .client ("logs" ).describe_log_streams .return_value = DEFAULT_LOG_STREAMS
13811381 boto_mock .client ("logs" ).get_log_events .side_effect = STREAM_LOG_EVENTS
1382- ims = sagemaker .Session (boto_session = boto_mock , sagemaker_client = Mock ())
1382+ ims = sagemaker .Session (boto_session = boto_mock , sagemaker_client = MagicMock ())
13831383 ims .sagemaker_client .describe_training_job .side_effect = [
13841384 IN_PROGRESS_DESCRIBE_JOB_RESULT ,
13851385 IN_PROGRESS_DESCRIBE_JOB_RESULT ,
@@ -1395,10 +1395,10 @@ def sagemaker_session_ready_lifecycle():
13951395
13961396@pytest .fixture ()
13971397def sagemaker_session_full_lifecycle ():
1398- boto_mock = Mock (name = "boto_session" )
1398+ boto_mock = MagicMock (name = "boto_session" )
13991399 boto_mock .client ("logs" ).describe_log_streams .side_effect = LIFECYCLE_LOG_STREAMS
14001400 boto_mock .client ("logs" ).get_log_events .side_effect = STREAM_LOG_EVENTS
1401- ims = sagemaker .Session (boto_session = boto_mock , sagemaker_client = Mock ())
1401+ ims = sagemaker .Session (boto_session = boto_mock , sagemaker_client = MagicMock ())
14021402 ims .sagemaker_client .describe_training_job .side_effect = [
14031403 IN_PROGRESS_DESCRIBE_JOB_RESULT ,
14041404 IN_PROGRESS_DESCRIBE_JOB_RESULT ,
0 commit comments