diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index e4373ab4dd..3f44bd93b5 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -632,7 +632,8 @@ def test_endpoint_from_production_variants_with_tags(sagemaker_session): Tags=tags) -def test_wait_for_tuning_job(sagemaker_session): +@patch('time.sleep') +def test_wait_for_tuning_job(sleep, sagemaker_session): hyperparameter_tuning_job_desc = {'HyperParameterTuningJobStatus': 'Completed'} sagemaker_session.sagemaker_client.describe_hyper_parameter_tuning_job = Mock( name='describe_hyper_parameter_tuning_job', return_value=hyperparameter_tuning_job_desc) @@ -661,7 +662,8 @@ def test_tune_job_status_none(sagemaker_session): assert result is None -def test_wait_for_transform_job_completed(sagemaker_session): +@patch('time.sleep') +def test_wait_for_transform_job_completed(sleep, sagemaker_session): transform_job_desc = {'TransformJobStatus': 'Completed'} sagemaker_session.sagemaker_client.describe_transform_job = Mock( name='describe_transform_job', return_value=transform_job_desc) @@ -669,7 +671,8 @@ def test_wait_for_transform_job_completed(sagemaker_session): assert sagemaker_session.wait_for_transform_job(JOB_NAME)['TransformJobStatus'] == 'Completed' -def test_wait_for_transform_job_in_progress(sagemaker_session): +@patch('time.sleep') +def test_wait_for_transform_job_in_progress(sleep, sagemaker_session): transform_job_desc_in_progress = {'TransformJobStatus': 'InProgress'} transform_job_desc_in_completed = {'TransformJobStatus': 'Completed'} sagemaker_session.sagemaker_client.describe_transform_job = Mock( diff --git a/tests/unit/test_tf_estimator.py b/tests/unit/test_tf_estimator.py index a2334dc736..0a155ce341 100644 --- a/tests/unit/test_tf_estimator.py +++ b/tests/unit/test_tf_estimator.py @@ -312,7 +312,8 @@ def test_run_tensorboard_locally_without_awscli_binary(time, strftime, popen, ca @patch('subprocess.Popen') @patch('time.strftime', return_value=TIMESTAMP) @patch('time.time', return_value=TIME) -def test_run_tensorboard_locally(time, strftime, popen, call, access, rmtree, mkdtemp, sync, sagemaker_session): +@patch('time.sleep') +def test_run_tensorboard_locally(sleep, time, strftime, popen, call, access, rmtree, mkdtemp, sync, sagemaker_session): tf = TensorFlow(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session, train_instance_count=INSTANCE_COUNT, train_instance_type=INSTANCE_TYPE) @@ -322,8 +323,7 @@ def test_run_tensorboard_locally(time, strftime, popen, call, access, rmtree, mk popen.assert_called_with(['tensorboard', '--logdir', '/my/temp/folder', '--host', 'localhost', '--port', '6006'], stderr=-1, - stdout=-1 - ) + stdout=-1) @patch('sagemaker.tensorflow.estimator.Tensorboard._sync_directories') @@ -335,7 +335,8 @@ def test_run_tensorboard_locally(time, strftime, popen, call, access, rmtree, mk @patch('subprocess.Popen') @patch('time.strftime', return_value=TIMESTAMP) @patch('time.time', return_value=TIME) -def test_run_tensorboard_locally_port_in_use(time, strftime, popen, call, access, socket, rmtree, mkdtemp, sync, +@patch('time.sleep') +def test_run_tensorboard_locally_port_in_use(sleep, time, strftime, popen, call, access, socket, rmtree, mkdtemp, sync, sagemaker_session): tf = TensorFlow(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session, train_instance_count=INSTANCE_COUNT, train_instance_type=INSTANCE_TYPE)