Skip to content

Commit 6ff1e23

Browse files
authored
Mock time.sleep in unit tests (#307)
Our unit tests currently take awhile to run. This change reduces the runtime from roughly 1 minute to 10 seconds.
1 parent cb4aa9d commit 6ff1e23

File tree

2 files changed

+11
-7
lines changed

2 files changed

+11
-7
lines changed

tests/unit/test_session.py

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -632,7 +632,8 @@ def test_endpoint_from_production_variants_with_tags(sagemaker_session):
632632
Tags=tags)
633633

634634

635-
def test_wait_for_tuning_job(sagemaker_session):
635+
@patch('time.sleep')
636+
def test_wait_for_tuning_job(sleep, sagemaker_session):
636637
hyperparameter_tuning_job_desc = {'HyperParameterTuningJobStatus': 'Completed'}
637638
sagemaker_session.sagemaker_client.describe_hyper_parameter_tuning_job = Mock(
638639
name='describe_hyper_parameter_tuning_job', return_value=hyperparameter_tuning_job_desc)
@@ -661,15 +662,17 @@ def test_tune_job_status_none(sagemaker_session):
661662
assert result is None
662663

663664

664-
def test_wait_for_transform_job_completed(sagemaker_session):
665+
@patch('time.sleep')
666+
def test_wait_for_transform_job_completed(sleep, sagemaker_session):
665667
transform_job_desc = {'TransformJobStatus': 'Completed'}
666668
sagemaker_session.sagemaker_client.describe_transform_job = Mock(
667669
name='describe_transform_job', return_value=transform_job_desc)
668670

669671
assert sagemaker_session.wait_for_transform_job(JOB_NAME)['TransformJobStatus'] == 'Completed'
670672

671673

672-
def test_wait_for_transform_job_in_progress(sagemaker_session):
674+
@patch('time.sleep')
675+
def test_wait_for_transform_job_in_progress(sleep, sagemaker_session):
673676
transform_job_desc_in_progress = {'TransformJobStatus': 'InProgress'}
674677
transform_job_desc_in_completed = {'TransformJobStatus': 'Completed'}
675678
sagemaker_session.sagemaker_client.describe_transform_job = Mock(

tests/unit/test_tf_estimator.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -312,7 +312,8 @@ def test_run_tensorboard_locally_without_awscli_binary(time, strftime, popen, ca
312312
@patch('subprocess.Popen')
313313
@patch('time.strftime', return_value=TIMESTAMP)
314314
@patch('time.time', return_value=TIME)
315-
def test_run_tensorboard_locally(time, strftime, popen, call, access, rmtree, mkdtemp, sync, sagemaker_session):
315+
@patch('time.sleep')
316+
def test_run_tensorboard_locally(sleep, time, strftime, popen, call, access, rmtree, mkdtemp, sync, sagemaker_session):
316317
tf = TensorFlow(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session,
317318
train_instance_count=INSTANCE_COUNT, train_instance_type=INSTANCE_TYPE)
318319

@@ -322,8 +323,7 @@ def test_run_tensorboard_locally(time, strftime, popen, call, access, rmtree, mk
322323

323324
popen.assert_called_with(['tensorboard', '--logdir', '/my/temp/folder', '--host', 'localhost', '--port', '6006'],
324325
stderr=-1,
325-
stdout=-1
326-
)
326+
stdout=-1)
327327

328328

329329
@patch('sagemaker.tensorflow.estimator.Tensorboard._sync_directories')
@@ -335,7 +335,8 @@ def test_run_tensorboard_locally(time, strftime, popen, call, access, rmtree, mk
335335
@patch('subprocess.Popen')
336336
@patch('time.strftime', return_value=TIMESTAMP)
337337
@patch('time.time', return_value=TIME)
338-
def test_run_tensorboard_locally_port_in_use(time, strftime, popen, call, access, socket, rmtree, mkdtemp, sync,
338+
@patch('time.sleep')
339+
def test_run_tensorboard_locally_port_in_use(sleep, time, strftime, popen, call, access, socket, rmtree, mkdtemp, sync,
339340
sagemaker_session):
340341
tf = TensorFlow(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session,
341342
train_instance_count=INSTANCE_COUNT, train_instance_type=INSTANCE_TYPE)

0 commit comments

Comments
 (0)