Skip to content

Commit a6f39de

Browse files
matthew29tangcopybara-github
matthew29tang
authored and
copybara-github
committed
chore: Speed up unit tests by setting mocked poll time to 0.05 seconds
PiperOrigin-RevId: 933822063
1 parent 13b685b commit a6f39de

18 files changed

Lines changed: 359 additions & 230 deletions

‎.kokoro/build.sh‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -43,15 +43,17 @@ export DOCKER_API_VERSION=1.39
4343
python3 -m pip install --upgrade --quiet uv nox
4444
python3 -m nox --version
4545

46-
# If this is a continuous build, send the test log to the FlakyBot.
46+
# Clean up the heavy .nox/ environment directory before artifact collection.
47+
# If this is a continuous build, also send the test log to the FlakyBot.
4748
# See https://github.com/googleapis/repo-automation-bots/tree/main/packages/flakybot.
48-
if [[ $KOKORO_BUILD_ARTIFACTS_SUBDIR = *"continuous"* ]]; then
49-
cleanup() {
49+
cleanup() {
50+
rm -rf .nox
51+
if [[ $KOKORO_BUILD_ARTIFACTS_SUBDIR = *"continuous"* ]]; then
5052
chmod +x $KOKORO_GFILE_DIR/linux_amd64/flakybot
5153
$KOKORO_GFILE_DIR/linux_amd64/flakybot
52-
}
53-
trap cleanup EXIT HUP
54-
fi
54+
fi
55+
}
56+
trap cleanup EXIT HUP
5557

5658
# If NOX_SESSION is set, it only runs the specified session,
5759
# otherwise run all the sessions.

‎noxfile.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -219,6 +219,13 @@ def default(session):
219219
"-n",
220220
"auto", # Use all available CPU cores
221221
"--quiet",
222+
"--durations=50",
223+
# Suppress redundant deprecation warnings to reduce JUnit XML size.
224+
"-W",
225+
"ignore:You are using a Python version (3.10.19):FutureWarning",
226+
# Disable capturing stdout/stderr for passed tests to reduce XML bloat.
227+
"-o",
228+
"junit_log_passing_tests=False",
222229
f"--junitxml=unit_{session.python}_sponge_log.xml",
223230
"--ignore=tests/unit/vertex_ray",
224231
"--ignore=tests/unit/vertex_adk",

‎tests/unit/aiplatform/test_autologging.py‎

Lines changed: 2 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -144,7 +144,7 @@
144144
_TEST_TF_EXPERIMENT_RUN_PARAMS = {
145145
"batch_size": "None",
146146
"class_weight": "None",
147-
"epochs": "5",
147+
"epochs": "1",
148148
"initial_epoch": "0",
149149
"max_queue_size": "10",
150150
"sample_weight": "None",
@@ -589,14 +589,6 @@ def build_and_train_test_tf_model():
589589
[1, 2],
590590
[2, 2],
591591
[2, 3],
592-
[1, 1],
593-
[1, 2],
594-
[2, 2],
595-
[2, 3],
596-
[1, 1],
597-
[1, 2],
598-
[2, 2],
599-
[2, 3],
600592
]
601593
)
602594
y = np.dot(X, np.array([1, 2])) + 3
@@ -605,7 +597,6 @@ def build_and_train_test_tf_model():
605597
[
606598
tf.keras.layers.Flatten(input_shape=(2,)),
607599
tf.keras.layers.Dense(128, activation="relu"),
608-
tf.keras.layers.Dropout(0.2),
609600
tf.keras.layers.Dense(1),
610601
]
611602
)
@@ -616,7 +607,7 @@ def build_and_train_test_tf_model():
616607
metrics=["accuracy"],
617608
)
618609

619-
model.fit(X, y, epochs=5)
610+
model.fit(X, y, epochs=1)
620611

621612

622613
@pytest.mark.usefixtures("google_auth_mock")

‎tests/unit/aiplatform/test_automl_forecasting_training_jobs.py‎

Lines changed: 30 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -303,12 +303,18 @@ class TestForecastingTrainingJob:
303303
def setup_method(self):
304304
importlib.reload(initializer)
305305
importlib.reload(aiplatform)
306+
self._job_wait_patcher = mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 0.05)
307+
self._log_wait_patcher = mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 0.05)
308+
self._job_wait_patcher.start()
309+
self._log_wait_patcher.start()
306310

307311
def teardown_method(self):
312+
self._job_wait_patcher.stop()
313+
self._log_wait_patcher.stop()
308314
initializer.global_pool.shutdown(wait=True)
309315

310-
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 1)
311-
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 1)
316+
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 0.05)
317+
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 0.05)
312318
@pytest.mark.parametrize("sync", [True, False])
313319
@pytest.mark.parametrize("training_job", _FORECASTING_JOB_MODEL_TYPES)
314320
def test_run_call_pipeline_service_create(
@@ -409,8 +415,8 @@ def test_run_call_pipeline_service_create(
409415

410416
assert job.state == gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED
411417

412-
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 1)
413-
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 1)
418+
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 0.05)
419+
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 0.05)
414420
@pytest.mark.parametrize("sync", [True, False])
415421
@pytest.mark.parametrize("training_job", _FORECASTING_JOB_MODEL_TYPES)
416422
def test_run_call_pipeline_service_create_with_timeout(
@@ -497,8 +503,8 @@ def test_run_call_pipeline_service_create_with_timeout(
497503
timeout=180.0,
498504
)
499505

500-
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 1)
501-
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 1)
506+
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 0.05)
507+
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 0.05)
502508
@pytest.mark.usefixtures("mock_pipeline_service_get")
503509
@pytest.mark.parametrize("sync", [True, False])
504510
@pytest.mark.parametrize("training_job", _FORECASTING_JOB_MODEL_TYPES)
@@ -579,8 +585,8 @@ def test_run_call_pipeline_if_no_model_display_name_nor_model_labels(
579585
timeout=None,
580586
)
581587

582-
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 1)
583-
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 1)
588+
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 0.05)
589+
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 0.05)
584590
@pytest.mark.usefixtures("mock_pipeline_service_get")
585591
@pytest.mark.parametrize("sync", [True, False])
586592
@pytest.mark.parametrize("training_job", _FORECASTING_JOB_MODEL_TYPES)
@@ -660,8 +666,8 @@ def test_run_call_pipeline_if_set_additional_experiments(
660666
timeout=None,
661667
)
662668

663-
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 1)
664-
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 1)
669+
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 0.05)
670+
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 0.05)
665671
@pytest.mark.usefixtures(
666672
"mock_pipeline_service_create",
667673
"mock_pipeline_service_get",
@@ -746,8 +752,8 @@ def test_run_called_twice_raises(
746752
holiday_regions=_TEST_TRAINING_HOLIDAY_REGIONS,
747753
)
748754

749-
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 1)
750-
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 1)
755+
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 0.05)
756+
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 0.05)
751757
@pytest.mark.parametrize("sync", [True, False])
752758
@pytest.mark.parametrize("training_job", _FORECASTING_JOB_MODEL_TYPES)
753759
def test_run_raises_if_pipeline_fails(
@@ -827,8 +833,8 @@ def test_raises_before_run_is_called(
827833
with pytest.raises(RuntimeError):
828834
job.state
829835

830-
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 1)
831-
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 1)
836+
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 0.05)
837+
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 0.05)
832838
@pytest.mark.parametrize("sync", [True, False])
833839
@pytest.mark.parametrize("training_job", _FORECASTING_JOB_MODEL_TYPES)
834840
def test_splits_fraction(
@@ -926,8 +932,8 @@ def test_splits_fraction(
926932
timeout=None,
927933
)
928934

929-
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 1)
930-
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 1)
935+
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 0.05)
936+
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 0.05)
931937
@pytest.mark.parametrize("sync", [True, False])
932938
@pytest.mark.parametrize("training_job", _FORECASTING_JOB_MODEL_TYPES)
933939
def test_splits_timestamp(
@@ -1027,8 +1033,8 @@ def test_splits_timestamp(
10271033
timeout=None,
10281034
)
10291035

1030-
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 1)
1031-
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 1)
1036+
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 0.05)
1037+
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 0.05)
10321038
@pytest.mark.parametrize("sync", [True, False])
10331039
@pytest.mark.parametrize("training_job", _FORECASTING_JOB_MODEL_TYPES)
10341040
def test_splits_predefined(
@@ -1122,8 +1128,8 @@ def test_splits_predefined(
11221128
timeout=None,
11231129
)
11241130

1125-
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 1)
1126-
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 1)
1131+
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 0.05)
1132+
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 0.05)
11271133
@pytest.mark.parametrize("sync", [True, False])
11281134
@pytest.mark.parametrize("training_job", _FORECASTING_JOB_MODEL_TYPES)
11291135
def test_splits_default(
@@ -1211,8 +1217,8 @@ def test_splits_default(
12111217
timeout=None,
12121218
)
12131219

1214-
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 1)
1215-
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 1)
1220+
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 0.05)
1221+
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 0.05)
12161222
@pytest.mark.usefixtures("mock_pipeline_service_get")
12171223
@pytest.mark.parametrize("sync", [True, False])
12181224
@pytest.mark.parametrize("training_job", _FORECASTING_JOB_MODEL_TYPES)
@@ -1294,8 +1300,8 @@ def test_run_call_pipeline_if_set_additional_experiments_probabilistic_inference
12941300
timeout=None,
12951301
)
12961302

1297-
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 1)
1298-
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 1)
1303+
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 0.05)
1304+
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 0.05)
12991305
@pytest.mark.usefixtures("mock_pipeline_service_get")
13001306
@pytest.mark.parametrize("sync", [True, False])
13011307
@pytest.mark.parametrize("training_job", _FORECASTING_JOB_MODEL_TYPES)

‎tests/unit/aiplatform/test_base.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -46,12 +46,12 @@ def _sync_object_with_future_result(self, result):
4646
@classmethod
4747
@base.optional_sync()
4848
def create(cls, x: int, sync=True) -> "_TestClass":
49-
time.sleep(1)
49+
time.sleep(0.05)
5050
return cls(x)
5151

5252
@base.optional_sync()
5353
def add(self, a: "_TestClass", sync=True) -> None:
54-
time.sleep(1)
54+
time.sleep(0.05)
5555
return self._add(a=a, sync=sync)
5656

5757
def _add(self, a: "_TestClass", sync=True) -> None:
@@ -63,14 +63,14 @@ class _TestClassDownStream(_TestClass):
6363
def add_and_create_new(
6464
self, a: Optional["_TestClass"] = None, sync=True
6565
) -> _TestClass:
66-
time.sleep(1)
66+
time.sleep(0.05)
6767
if a:
6868
return _TestClass(self.x + a.x)
6969
return None
7070

7171
@base.optional_sync(return_input_arg="a", bind_future_to_self=False)
7272
def add_to_input_arg(self, a: "_TestClass", sync=True) -> _TestClass:
73-
time.sleep(1)
73+
time.sleep(0.05)
7474
a._add(self)
7575
return a
7676

‎tests/unit/aiplatform/test_batch_prediction_job_preview.py‎

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -123,10 +123,23 @@ def create_batch_prediction_job_mock():
123123
class TestBatchPredictionJobPreview:
124124

125125
def setup_method(self):
126+
from google.cloud.aiplatform import jobs
126127
reload(initializer)
127128
reload(aiplatform)
129+
self._preview_job_wait_patcher = mock.patch.object(preview_jobs, "_JOB_WAIT_TIME", 0.05)
130+
self._preview_log_wait_patcher = mock.patch.object(preview_jobs, "_LOG_WAIT_TIME", 0.05)
131+
self._job_wait_patcher = mock.patch.object(jobs, "_JOB_WAIT_TIME", 0.05)
132+
self._log_wait_patcher = mock.patch.object(jobs, "_LOG_WAIT_TIME", 0.05)
133+
self._preview_job_wait_patcher.start()
134+
self._preview_log_wait_patcher.start()
135+
self._job_wait_patcher.start()
136+
self._log_wait_patcher.start()
128137

129138
def teardown_method(self):
139+
self._preview_job_wait_patcher.stop()
140+
self._preview_log_wait_patcher.stop()
141+
self._job_wait_patcher.stop()
142+
self._log_wait_patcher.stop()
130143
initializer.global_pool.shutdown(wait=True)
131144

132145
def test_init_batch_prediction_job(self, get_batch_prediction_job_mock):
@@ -160,8 +173,8 @@ def test_batch_prediction_job_done_get(self, get_batch_prediction_job_mock):
160173
assert bp.done() is False
161174
assert get_batch_prediction_job_mock.call_count == 2
162175

163-
@mock.patch.object(preview_jobs, "_JOB_WAIT_TIME", 1)
164-
@mock.patch.object(preview_jobs, "_LOG_WAIT_TIME", 1)
176+
@mock.patch.object(preview_jobs, "_JOB_WAIT_TIME", 0.05)
177+
@mock.patch.object(preview_jobs, "_LOG_WAIT_TIME", 0.05)
165178
@pytest.mark.parametrize("sync", [True, False])
166179
@pytest.mark.usefixtures("get_batch_prediction_job_mock")
167180
def test_batch_predict_create_with_reservation(
@@ -228,8 +241,8 @@ def test_batch_predict_create_with_reservation(
228241
timeout=None,
229242
)
230243

231-
@mock.patch.object(preview_jobs, "_JOB_WAIT_TIME", 1)
232-
@mock.patch.object(preview_jobs, "_LOG_WAIT_TIME", 1)
244+
@mock.patch.object(preview_jobs, "_JOB_WAIT_TIME", 0.05)
245+
@mock.patch.object(preview_jobs, "_LOG_WAIT_TIME", 0.05)
233246
@pytest.mark.usefixtures("get_batch_prediction_job_mock")
234247
def test_batch_predict_job_submit(self, create_batch_prediction_job_mock):
235248
aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)

‎tests/unit/aiplatform/test_custom_job.py‎

Lines changed: 16 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -633,8 +633,14 @@ class TestCustomJob:
633633
def setup_method(self):
634634
reload(aiplatform.initializer)
635635
reload(aiplatform)
636+
self._job_wait_patcher = mock.patch.object(jobs, "_JOB_WAIT_TIME", 0.05)
637+
self._log_wait_patcher = mock.patch.object(jobs, "_LOG_WAIT_TIME", 0.05)
638+
self._job_wait_patcher.start()
639+
self._log_wait_patcher.start()
636640

637641
def teardown_method(self):
642+
self._job_wait_patcher.stop()
643+
self._log_wait_patcher.stop()
638644
aiplatform.initializer.global_pool.shutdown(wait=True)
639645

640646
@pytest.mark.parametrize("sync", [True, False])
@@ -874,8 +880,8 @@ def test_submit_custom_job_with_experiments(
874880
)
875881

876882
@pytest.mark.parametrize("sync", [True, False])
877-
@mock.patch.object(jobs, "_JOB_WAIT_TIME", 1)
878-
@mock.patch.object(jobs, "_LOG_WAIT_TIME", 1)
883+
@mock.patch.object(jobs, "_JOB_WAIT_TIME", 0.05)
884+
@mock.patch.object(jobs, "_LOG_WAIT_TIME", 0.05)
879885
def test_create_custom_job_with_timeout(
880886
self, create_custom_job_mock, get_custom_job_mock, sync
881887
):
@@ -1317,8 +1323,8 @@ def test_create_from_local_script_raises_with_no_staging_bucket(
13171323
"update_context_mock",
13181324
)
13191325
@pytest.mark.parametrize("sync", [True, False])
1320-
@mock.patch.object(jobs, "_JOB_WAIT_TIME", 1)
1321-
@mock.patch.object(jobs, "_LOG_WAIT_TIME", 1)
1326+
@mock.patch.object(jobs, "_JOB_WAIT_TIME", 0.05)
1327+
@mock.patch.object(jobs, "_LOG_WAIT_TIME", 0.05)
13221328
def test_create_from_local_script_prebuilt_container_with_all_args(
13231329
self, get_custom_job_mock, create_custom_job_mock, sync
13241330
):
@@ -1381,8 +1387,8 @@ def test_create_from_local_script_prebuilt_container_with_all_args(
13811387
"update_context_mock",
13821388
)
13831389
@pytest.mark.parametrize("sync", [True, False])
1384-
@mock.patch.object(jobs, "_JOB_WAIT_TIME", 1)
1385-
@mock.patch.object(jobs, "_LOG_WAIT_TIME", 1)
1390+
@mock.patch.object(jobs, "_JOB_WAIT_TIME", 0.05)
1391+
@mock.patch.object(jobs, "_LOG_WAIT_TIME", 0.05)
13861392
def test_create_from_local_script_custom_container_with_all_args(
13871393
self, get_custom_job_mock, create_custom_job_mock, sync
13881394
):
@@ -1458,8 +1464,8 @@ def test_create_from_local_script_enable_autolog_no_experiment_error(self):
14581464
job.run()
14591465

14601466
@pytest.mark.parametrize("sync", [True, False])
1461-
@mock.patch.object(jobs, "_JOB_WAIT_TIME", 1)
1462-
@mock.patch.object(jobs, "_LOG_WAIT_TIME", 1)
1467+
@mock.patch.object(jobs, "_JOB_WAIT_TIME", 0.05)
1468+
@mock.patch.object(jobs, "_LOG_WAIT_TIME", 0.05)
14631469
def test_create_custom_job_with_enable_web_access(
14641470
self,
14651471
create_custom_job_mock_with_enable_web_access,
@@ -1525,8 +1531,8 @@ def test_get_web_access_uris(self, get_custom_job_mock_with_enable_web_access):
15251531
assert job.web_access_uris == _TEST_WEB_ACCESS_URIS
15261532
break
15271533

1528-
@mock.patch.object(jobs, "_JOB_WAIT_TIME", 1)
1529-
@mock.patch.object(jobs, "_LOG_WAIT_TIME", 1)
1534+
@mock.patch.object(jobs, "_JOB_WAIT_TIME", 0.05)
1535+
@mock.patch.object(jobs, "_LOG_WAIT_TIME", 0.05)
15301536
def test_log_access_web_uris_after_get(
15311537
self, get_custom_job_mock_with_enable_web_access
15321538
):

‎tests/unit/aiplatform/test_custom_job_persistent_resource.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -152,8 +152,14 @@ class TestCustomJobPersistentResource:
152152
def setup_method(self):
153153
reload(aiplatform.initializer)
154154
reload(aiplatform)
155+
self._job_wait_patcher = mock.patch.object(jobs, "_JOB_WAIT_TIME", 0.05)
156+
self._log_wait_patcher = mock.patch.object(jobs, "_LOG_WAIT_TIME", 0.05)
157+
self._job_wait_patcher.start()
158+
self._log_wait_patcher.start()
155159

156160
def teardown_method(self):
161+
self._job_wait_patcher.stop()
162+
self._log_wait_patcher.stop()
157163
aiplatform.initializer.global_pool.shutdown(wait=True)
158164

159165
@pytest.mark.parametrize("sync", [True, False])

0 commit comments

Comments
 (0)