From 0531a5e5a84b0b283fe5c3bfb88c9461792cf2d7 Mon Sep 17 00:00:00 2001 From: Lauren Yu <6631887+laurenyu@users.noreply.github.com> Date: Thu, 19 Jul 2018 10:59:30 -0700 Subject: [PATCH] Mock time.sleep in unit tests --- tests/unit/test_session.py | 9 ++++++--- tests/unit/test_tf_estimator.py | 9 +++++---- 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index 00a6fd82be..2a92965115 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -592,7 +592,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) @@ -621,7 +622,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) @@ -629,7 +631,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)