Skip to content

chore: Speed up unit tests by setting mocked poll time to 0.05 seconds #6889

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
copybara-service merged 1 commit into from
Jun 17, 2026
Merged

chore: Speed up unit tests by setting mocked poll time to 0.05 seconds #6889

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 8 additions & 6 deletions .kokoro/build.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,15 +43,17 @@ export DOCKER_API_VERSION=1.39
python3 -m pip install --upgrade --quiet uv nox
python3 -m nox --version

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

# If NOX_SESSION is set, it only runs the specified session,
# otherwise run all the sessions.
Expand Down
7 changes: 7 additions & 0 deletions noxfile.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -219,6 +219,13 @@ def default(session):
"-n",
"auto", # Use all available CPU cores
"--quiet",
"--durations=50",
# Suppress redundant deprecation warnings to reduce JUnit XML size.
"-W",
"ignore:You are using a Python version (3.10.19):FutureWarning",
# Disable capturing stdout/stderr for passed tests to reduce XML bloat.
"-o",
"junit_log_passing_tests=False",
f"--junitxml=unit_{session.python}_sponge_log.xml",
"--ignore=tests/unit/vertex_ray",
"--ignore=tests/unit/vertex_adk",
Expand Down
13 changes: 2 additions & 11 deletions tests/unit/aiplatform/test_autologging.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -144,7 +144,7 @@
_TEST_TF_EXPERIMENT_RUN_PARAMS = {
"batch_size": "None",
"class_weight": "None",
"epochs": "5",
"epochs": "1",
"initial_epoch": "0",
"max_queue_size": "10",
"sample_weight": "None",
Expand Down Expand Up @@ -589,14 +589,6 @@ def build_and_train_test_tf_model():
[1, 2],
[2, 2],
[2, 3],
[1, 1],
[1, 2],
[2, 2],
[2, 3],
[1, 1],
[1, 2],
[2, 2],
[2, 3],
]
)
y = np.dot(X, np.array([1, 2])) + 3
Expand All @@ -605,7 +597,6 @@ def build_and_train_test_tf_model():
[
tf.keras.layers.Flatten(input_shape=(2,)),
tf.keras.layers.Dense(128, activation="relu"),
tf.keras.layers.Dropout(0.2),
tf.keras.layers.Dense(1),
]
)
Expand All @@ -616,7 +607,7 @@ def build_and_train_test_tf_model():
metrics=["accuracy"],
)

model.fit(X, y, epochs=5)
model.fit(X, y, epochs=1)


@pytest.mark.usefixtures("google_auth_mock")
Expand Down
54 changes: 30 additions & 24 deletions tests/unit/aiplatform/test_automl_forecasting_training_jobs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -303,12 +303,18 @@ class TestForecastingTrainingJob:
def setup_method(self):
importlib.reload(initializer)
importlib.reload(aiplatform)
self._job_wait_patcher = mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 0.05)
self._log_wait_patcher = mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 0.05)
self._job_wait_patcher.start()
self._log_wait_patcher.start()

def teardown_method(self):
self._job_wait_patcher.stop()
self._log_wait_patcher.stop()
initializer.global_pool.shutdown(wait=True)

@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 1)
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 1)
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 0.05)
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 0.05)
@pytest.mark.parametrize("sync", [True, False])
@pytest.mark.parametrize("training_job", _FORECASTING_JOB_MODEL_TYPES)
def test_run_call_pipeline_service_create(
Expand Down Expand Up @@ -409,8 +415,8 @@ def test_run_call_pipeline_service_create(

assert job.state == gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED

@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 1)
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 1)
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 0.05)
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 0.05)
@pytest.mark.parametrize("sync", [True, False])
@pytest.mark.parametrize("training_job", _FORECASTING_JOB_MODEL_TYPES)
def test_run_call_pipeline_service_create_with_timeout(
Expand Down Expand Up @@ -497,8 +503,8 @@ def test_run_call_pipeline_service_create_with_timeout(
timeout=180.0,
)

@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 1)
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 1)
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 0.05)
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 0.05)
@pytest.mark.usefixtures("mock_pipeline_service_get")
@pytest.mark.parametrize("sync", [True, False])
@pytest.mark.parametrize("training_job", _FORECASTING_JOB_MODEL_TYPES)
Expand Down Expand Up @@ -579,8 +585,8 @@ def test_run_call_pipeline_if_no_model_display_name_nor_model_labels(
timeout=None,
)

@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 1)
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 1)
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 0.05)
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 0.05)
@pytest.mark.usefixtures("mock_pipeline_service_get")
@pytest.mark.parametrize("sync", [True, False])
@pytest.mark.parametrize("training_job", _FORECASTING_JOB_MODEL_TYPES)
Expand Down Expand Up @@ -660,8 +666,8 @@ def test_run_call_pipeline_if_set_additional_experiments(
timeout=None,
)

@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 1)
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 1)
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 0.05)
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 0.05)
@pytest.mark.usefixtures(
"mock_pipeline_service_create",
"mock_pipeline_service_get",
Expand Down Expand Up @@ -746,8 +752,8 @@ def test_run_called_twice_raises(
holiday_regions=_TEST_TRAINING_HOLIDAY_REGIONS,
)

@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 1)
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 1)
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 0.05)
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 0.05)
@pytest.mark.parametrize("sync", [True, False])
@pytest.mark.parametrize("training_job", _FORECASTING_JOB_MODEL_TYPES)
def test_run_raises_if_pipeline_fails(
Expand Down Expand Up @@ -827,8 +833,8 @@ def test_raises_before_run_is_called(
with pytest.raises(RuntimeError):
job.state

@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 1)
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 1)
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 0.05)
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 0.05)
@pytest.mark.parametrize("sync", [True, False])
@pytest.mark.parametrize("training_job", _FORECASTING_JOB_MODEL_TYPES)
def test_splits_fraction(
Expand Down Expand Up @@ -926,8 +932,8 @@ def test_splits_fraction(
timeout=None,
)

@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 1)
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 1)
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 0.05)
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 0.05)
@pytest.mark.parametrize("sync", [True, False])
@pytest.mark.parametrize("training_job", _FORECASTING_JOB_MODEL_TYPES)
def test_splits_timestamp(
Expand Down Expand Up @@ -1027,8 +1033,8 @@ def test_splits_timestamp(
timeout=None,
)

@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 1)
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 1)
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 0.05)
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 0.05)
@pytest.mark.parametrize("sync", [True, False])
@pytest.mark.parametrize("training_job", _FORECASTING_JOB_MODEL_TYPES)
def test_splits_predefined(
Expand Down Expand Up @@ -1122,8 +1128,8 @@ def test_splits_predefined(
timeout=None,
)

@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 1)
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 1)
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 0.05)
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 0.05)
@pytest.mark.parametrize("sync", [True, False])
@pytest.mark.parametrize("training_job", _FORECASTING_JOB_MODEL_TYPES)
def test_splits_default(
Expand Down Expand Up @@ -1211,8 +1217,8 @@ def test_splits_default(
timeout=None,
)

@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 1)
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 1)
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 0.05)
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 0.05)
@pytest.mark.usefixtures("mock_pipeline_service_get")
@pytest.mark.parametrize("sync", [True, False])
@pytest.mark.parametrize("training_job", _FORECASTING_JOB_MODEL_TYPES)
Expand Down Expand Up @@ -1294,8 +1300,8 @@ def test_run_call_pipeline_if_set_additional_experiments_probabilistic_inference
timeout=None,
)

@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 1)
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 1)
@mock.patch.object(training_jobs, "_JOB_WAIT_TIME", 0.05)
@mock.patch.object(training_jobs, "_LOG_WAIT_TIME", 0.05)
@pytest.mark.usefixtures("mock_pipeline_service_get")
@pytest.mark.parametrize("sync", [True, False])
@pytest.mark.parametrize("training_job", _FORECASTING_JOB_MODEL_TYPES)
Expand Down
8 changes: 4 additions & 4 deletions tests/unit/aiplatform/test_base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,12 +46,12 @@ def _sync_object_with_future_result(self, result):
@classmethod
@base.optional_sync()
def create(cls, x: int, sync=True) -> "_TestClass":
time.sleep(1)
time.sleep(0.05)
return cls(x)

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

def _add(self, a: "_TestClass", sync=True) -> None:
Expand All @@ -63,14 +63,14 @@ class _TestClassDownStream(_TestClass):
def add_and_create_new(
self, a: Optional["_TestClass"] = None, sync=True
) -> _TestClass:
time.sleep(1)
time.sleep(0.05)
if a:
return _TestClass(self.x + a.x)
return None

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

Expand Down
21 changes: 17 additions & 4 deletions tests/unit/aiplatform/test_batch_prediction_job_preview.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -123,10 +123,23 @@ def create_batch_prediction_job_mock():
class TestBatchPredictionJobPreview:

def setup_method(self):
from google.cloud.aiplatform import jobs
reload(initializer)
reload(aiplatform)
self._preview_job_wait_patcher = mock.patch.object(preview_jobs, "_JOB_WAIT_TIME", 0.05)
self._preview_log_wait_patcher = mock.patch.object(preview_jobs, "_LOG_WAIT_TIME", 0.05)
self._job_wait_patcher = mock.patch.object(jobs, "_JOB_WAIT_TIME", 0.05)
self._log_wait_patcher = mock.patch.object(jobs, "_LOG_WAIT_TIME", 0.05)
self._preview_job_wait_patcher.start()
self._preview_log_wait_patcher.start()
self._job_wait_patcher.start()
self._log_wait_patcher.start()

def teardown_method(self):
self._preview_job_wait_patcher.stop()
self._preview_log_wait_patcher.stop()
self._job_wait_patcher.stop()
self._log_wait_patcher.stop()
initializer.global_pool.shutdown(wait=True)

def test_init_batch_prediction_job(self, get_batch_prediction_job_mock):
Expand Down Expand Up @@ -160,8 +173,8 @@ def test_batch_prediction_job_done_get(self, get_batch_prediction_job_mock):
assert bp.done() is False
assert get_batch_prediction_job_mock.call_count == 2

@mock.patch.object(preview_jobs, "_JOB_WAIT_TIME", 1)
@mock.patch.object(preview_jobs, "_LOG_WAIT_TIME", 1)
@mock.patch.object(preview_jobs, "_JOB_WAIT_TIME", 0.05)
@mock.patch.object(preview_jobs, "_LOG_WAIT_TIME", 0.05)
@pytest.mark.parametrize("sync", [True, False])
@pytest.mark.usefixtures("get_batch_prediction_job_mock")
def test_batch_predict_create_with_reservation(
Expand Down Expand Up @@ -228,8 +241,8 @@ def test_batch_predict_create_with_reservation(
timeout=None,
)

@mock.patch.object(preview_jobs, "_JOB_WAIT_TIME", 1)
@mock.patch.object(preview_jobs, "_LOG_WAIT_TIME", 1)
@mock.patch.object(preview_jobs, "_JOB_WAIT_TIME", 0.05)
@mock.patch.object(preview_jobs, "_LOG_WAIT_TIME", 0.05)
@pytest.mark.usefixtures("get_batch_prediction_job_mock")
def test_batch_predict_job_submit(self, create_batch_prediction_job_mock):
aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
Expand Down
26 changes: 16 additions & 10 deletions tests/unit/aiplatform/test_custom_job.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -633,8 +633,14 @@ class TestCustomJob:
def setup_method(self):
reload(aiplatform.initializer)
reload(aiplatform)
self._job_wait_patcher = mock.patch.object(jobs, "_JOB_WAIT_TIME", 0.05)
self._log_wait_patcher = mock.patch.object(jobs, "_LOG_WAIT_TIME", 0.05)
self._job_wait_patcher.start()
self._log_wait_patcher.start()

def teardown_method(self):
self._job_wait_patcher.stop()
self._log_wait_patcher.stop()
aiplatform.initializer.global_pool.shutdown(wait=True)

@pytest.mark.parametrize("sync", [True, False])
Expand Down Expand Up @@ -874,8 +880,8 @@ def test_submit_custom_job_with_experiments(
)

@pytest.mark.parametrize("sync", [True, False])
@mock.patch.object(jobs, "_JOB_WAIT_TIME", 1)
@mock.patch.object(jobs, "_LOG_WAIT_TIME", 1)
@mock.patch.object(jobs, "_JOB_WAIT_TIME", 0.05)
@mock.patch.object(jobs, "_LOG_WAIT_TIME", 0.05)
def test_create_custom_job_with_timeout(
self, create_custom_job_mock, get_custom_job_mock, sync
):
Expand Down Expand Up @@ -1317,8 +1323,8 @@ def test_create_from_local_script_raises_with_no_staging_bucket(
"update_context_mock",
)
@pytest.mark.parametrize("sync", [True, False])
@mock.patch.object(jobs, "_JOB_WAIT_TIME", 1)
@mock.patch.object(jobs, "_LOG_WAIT_TIME", 1)
@mock.patch.object(jobs, "_JOB_WAIT_TIME", 0.05)
@mock.patch.object(jobs, "_LOG_WAIT_TIME", 0.05)
def test_create_from_local_script_prebuilt_container_with_all_args(
self, get_custom_job_mock, create_custom_job_mock, sync
):
Expand Down Expand Up @@ -1381,8 +1387,8 @@ def test_create_from_local_script_prebuilt_container_with_all_args(
"update_context_mock",
)
@pytest.mark.parametrize("sync", [True, False])
@mock.patch.object(jobs, "_JOB_WAIT_TIME", 1)
@mock.patch.object(jobs, "_LOG_WAIT_TIME", 1)
@mock.patch.object(jobs, "_JOB_WAIT_TIME", 0.05)
@mock.patch.object(jobs, "_LOG_WAIT_TIME", 0.05)
def test_create_from_local_script_custom_container_with_all_args(
self, get_custom_job_mock, create_custom_job_mock, sync
):
Expand Down Expand Up @@ -1458,8 +1464,8 @@ def test_create_from_local_script_enable_autolog_no_experiment_error(self):
job.run()

@pytest.mark.parametrize("sync", [True, False])
@mock.patch.object(jobs, "_JOB_WAIT_TIME", 1)
@mock.patch.object(jobs, "_LOG_WAIT_TIME", 1)
@mock.patch.object(jobs, "_JOB_WAIT_TIME", 0.05)
@mock.patch.object(jobs, "_LOG_WAIT_TIME", 0.05)
def test_create_custom_job_with_enable_web_access(
self,
create_custom_job_mock_with_enable_web_access,
Expand Down Expand Up @@ -1525,8 +1531,8 @@ def test_get_web_access_uris(self, get_custom_job_mock_with_enable_web_access):
assert job.web_access_uris == _TEST_WEB_ACCESS_URIS
break

@mock.patch.object(jobs, "_JOB_WAIT_TIME", 1)
@mock.patch.object(jobs, "_LOG_WAIT_TIME", 1)
@mock.patch.object(jobs, "_JOB_WAIT_TIME", 0.05)
@mock.patch.object(jobs, "_LOG_WAIT_TIME", 0.05)
def test_log_access_web_uris_after_get(
self, get_custom_job_mock_with_enable_web_access
):
Expand Down
6 changes: 6 additions & 0 deletions tests/unit/aiplatform/test_custom_job_persistent_resource.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -152,8 +152,14 @@ class TestCustomJobPersistentResource:
def setup_method(self):
reload(aiplatform.initializer)
reload(aiplatform)
self._job_wait_patcher = mock.patch.object(jobs, "_JOB_WAIT_TIME", 0.05)
self._log_wait_patcher = mock.patch.object(jobs, "_LOG_WAIT_TIME", 0.05)
self._job_wait_patcher.start()
self._log_wait_patcher.start()

def teardown_method(self):
self._job_wait_patcher.stop()
self._log_wait_patcher.stop()
aiplatform.initializer.global_pool.shutdown(wait=True)

@pytest.mark.parametrize("sync", [True, False])
Expand Down
Loading
Loading