Skip to content

Commit b989dbb

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
chore: separating tests for efficiency
PiperOrigin-RevId: 524403464
1 parent 8e91a58 commit b989dbb

10 files changed

+266
-180
lines changed

‎tests/unit/aiplatform/constants.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,15 @@ class TrainingJobConstants:
8383
_TEST_REDUCTION_SERVER_CONTAINER_URI = (
8484
"us-docker.pkg.dev/vertex-ai-restricted/training/reductionserver:latest"
8585
)
86+
_TEST_DATASET_DISPLAY_NAME = "test-dataset-display-name"
87+
_TEST_DATASET_NAME = "test-dataset-name"
88+
_TEST_DISPLAY_NAME = "test-display-name"
89+
_TEST_BUCKET_NAME = "test-bucket"
90+
_TEST_GCS_PATH_WITHOUT_BUCKET = "path/to/folder"
91+
_TEST_GCS_PATH = f"{_TEST_BUCKET_NAME}/{_TEST_GCS_PATH_WITHOUT_BUCKET}"
92+
_TEST_GCS_PATH_WITH_TRAILING_SLASH = f"{_TEST_GCS_PATH}/"
93+
_TEST_MODEL_DISPLAY_NAME = "model-display-name"
94+
_TEST_MODEL_LABELS = {"model_key": "model_value"}
8695
_TEST_STAGING_BUCKET = "gs://test-staging-bucket"
8796
_TEST_DISPLAY_NAME = "my_job_1234"
8897
_TEST_BASE_OUTPUT_DIR = f"{_TEST_STAGING_BUCKET}/{_TEST_DISPLAY_NAME}"
@@ -146,6 +155,11 @@ class TrainingJobConstants:
146155
_TEST_TEST_FRACTION_SPLIT = 0.2
147156
_TEST_BOOT_DISK_TYPE_DEFAULT = "pd-ssd"
148157
_TEST_BOOT_DISK_SIZE_GB_DEFAULT = 100
158+
# # DUPLICATE: THIS NEEDS TO BE MOVED TO THE TRAINING JOB TEST THAT USES IT
159+
_TEST_PIPELINE_RESOURCE_NAME = (
160+
"projects/my-project/locations/us-central1/trainingPipelines/12345"
161+
)
162+
_TEST_DEFAULT_ENCRYPTION_KEY_NAME = "key_default"
149163

150164

151165
@dataclasses.dataclass(frozen=True)

‎tests/unit/aiplatform/test_automl_forecasting_training_jobs.py

Lines changed: 26 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -40,17 +40,26 @@
4040
)
4141
from google.protobuf import json_format
4242
from google.protobuf import struct_pb2
43+
import constants as test_constants
4344

44-
_TEST_BUCKET_NAME = "test-bucket"
45-
_TEST_GCS_PATH_WITHOUT_BUCKET = "path/to/folder"
46-
_TEST_GCS_PATH = f"{_TEST_BUCKET_NAME}/{_TEST_GCS_PATH_WITHOUT_BUCKET}"
47-
_TEST_GCS_PATH_WITH_TRAILING_SLASH = f"{_TEST_GCS_PATH}/"
48-
_TEST_PROJECT = "test-project"
45+
_TEST_BUCKET_NAME = test_constants.TrainingJobConstants._TEST_BUCKET_NAME
46+
_TEST_GCS_PATH_WITHOUT_BUCKET = (
47+
test_constants.TrainingJobConstants._TEST_GCS_PATH_WITHOUT_BUCKET
48+
)
49+
_TEST_GCS_PATH = test_constants.TrainingJobConstants._TEST_GCS_PATH
50+
_TEST_GCS_PATH_WITH_TRAILING_SLASH = (
51+
test_constants.TrainingJobConstants._TEST_GCS_PATH_WITH_TRAILING_SLASH
52+
)
53+
_TEST_PROJECT = test_constants.ProjectConstants._TEST_PROJECT
4954

50-
_TEST_DATASET_DISPLAY_NAME = "test-dataset-display-name"
51-
_TEST_DATASET_NAME = "test-dataset-name"
52-
_TEST_DISPLAY_NAME = "test-display-name"
53-
_TEST_TRAINING_CONTAINER_IMAGE = "gcr.io/test-training/container:image"
55+
_TEST_DATASET_DISPLAY_NAME = (
56+
test_constants.TrainingJobConstants._TEST_DATASET_DISPLAY_NAME
57+
)
58+
_TEST_DATASET_NAME = test_constants.TrainingJobConstants._TEST_DATASET_NAME
59+
_TEST_DISPLAY_NAME = test_constants.TrainingJobConstants._TEST_DISPLAY_NAME
60+
_TEST_TRAINING_CONTAINER_IMAGE = (
61+
test_constants.TrainingJobConstants._TEST_TRAINING_CONTAINER_IMAGE
62+
)
5463
_TEST_METADATA_SCHEMA_URI_TIMESERIES = schema.dataset.metadata.time_series
5564
_TEST_METADATA_SCHEMA_URI_NONTIMESERIES = schema.dataset.metadata.image
5665

@@ -155,23 +164,25 @@
155164
struct_pb2.Value(),
156165
)
157166

158-
_TEST_DATASET_NAME = "test-dataset-name"
167+
_TEST_DATASET_NAME = test_constants.TrainingJobConstants._TEST_DATASET_NAME
159168

160-
_TEST_MODEL_DISPLAY_NAME = "model-display-name"
169+
_TEST_MODEL_DISPLAY_NAME = test_constants.TrainingJobConstants._TEST_MODEL_DISPLAY_NAME
161170

162-
_TEST_LABELS = {"key": "value"}
163-
_TEST_MODEL_LABELS = {"model_key": "model_value"}
171+
_TEST_LABELS = test_constants.ProjectConstants._TEST_LABELS
172+
_TEST_MODEL_LABELS = test_constants.TrainingJobConstants._TEST_MODEL_LABELS
164173

165174
_TEST_PREDEFINED_SPLIT_COLUMN_NAME = "split"
166175

167176
_TEST_MODEL_NAME = "projects/my-project/locations/us-central1/models/12345"
168177

169178
_TEST_PIPELINE_RESOURCE_NAME = (
170-
"projects/my-project/locations/us-central1/trainingPipelines/12345"
179+
test_constants.TrainingJobConstants._TEST_PIPELINE_RESOURCE_NAME
171180
)
172181

173182
# CMEK encryption
174-
_TEST_DEFAULT_ENCRYPTION_KEY_NAME = "key_default"
183+
_TEST_DEFAULT_ENCRYPTION_KEY_NAME = (
184+
test_constants.TrainingJobConstants._TEST_DEFAULT_ENCRYPTION_KEY_NAME
185+
)
175186
_TEST_DEFAULT_ENCRYPTION_SPEC = gca_encryption_spec.EncryptionSpec(
176187
kms_key_name=_TEST_DEFAULT_ENCRYPTION_KEY_NAME
177188
)

‎tests/unit/aiplatform/test_automl_image_training_jobs.py

Lines changed: 15 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -41,12 +41,15 @@
4141
pipeline_state as gca_pipeline_state,
4242
training_pipeline as gca_training_pipeline,
4343
)
44+
import constants as test_constants
4445

45-
_TEST_PROJECT = "test-project"
46-
_TEST_LOCATION = "us-central1"
47-
_TEST_DATASET_DISPLAY_NAME = "test-dataset-display-name"
48-
_TEST_DATASET_NAME = "test-dataset-name"
49-
_TEST_DISPLAY_NAME = "test-display-name"
46+
_TEST_PROJECT = test_constants.ProjectConstants._TEST_PROJECT
47+
_TEST_LOCATION = test_constants.ProjectConstants._TEST_LOCATION
48+
_TEST_DATASET_DISPLAY_NAME = (
49+
test_constants.TrainingJobConstants._TEST_DATASET_DISPLAY_NAME
50+
)
51+
_TEST_DATASET_NAME = test_constants.TrainingJobConstants._TEST_DATASET_NAME
52+
_TEST_DISPLAY_NAME = test_constants.TrainingJobConstants._TEST_DISPLAY_NAME
5053
_TEST_METADATA_SCHEMA_URI_IMAGE = schema.dataset.metadata.image
5154

5255
_TEST_TRAINING_BUDGET_MILLI_NODE_HOURS = 7500
@@ -57,12 +60,11 @@
5760
_TEST_PREDICTION_TYPE_ICN = "classification"
5861
_TEST_PREDICTION_TYPE_IOD = "object_detection"
5962

60-
_TEST_DATASET_NAME = "test-dataset-name"
61-
_TEST_MODEL_DISPLAY_NAME = "model-display-name"
63+
_TEST_MODEL_DISPLAY_NAME = test_constants.TrainingJobConstants._TEST_MODEL_DISPLAY_NAME
6264
_TEST_MODEL_ID = "98777645321"
6365

64-
_TEST_LABELS = {"key": "value"}
65-
_TEST_MODEL_LABELS = {"model_key": "model_value"}
66+
_TEST_LABELS = test_constants.ProjectConstants._TEST_LABELS
67+
_TEST_MODEL_LABELS = test_constants.TrainingJobConstants._TEST_MODEL_LABELS
6668

6769
_TEST_TRAINING_TASK_INPUTS = json_format.ParseDict(
6870
{
@@ -109,11 +111,13 @@
109111
)
110112

111113
_TEST_PIPELINE_RESOURCE_NAME = (
112-
f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/trainingPipelines/12345"
114+
test_constants.TrainingJobConstants._TEST_PIPELINE_RESOURCE_NAME
113115
)
114116

115117
# CMEK encryption
116-
_TEST_DEFAULT_ENCRYPTION_KEY_NAME = "key_default"
118+
_TEST_DEFAULT_ENCRYPTION_KEY_NAME = (
119+
test_constants.TrainingJobConstants._TEST_DEFAULT_ENCRYPTION_KEY_NAME
120+
)
117121
_TEST_DEFAULT_ENCRYPTION_SPEC = gca_encryption_spec.EncryptionSpec(
118122
kms_key_name=_TEST_DEFAULT_ENCRYPTION_KEY_NAME
119123
)

‎tests/unit/aiplatform/test_automl_tabular_training_jobs.py

Lines changed: 22 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -40,16 +40,23 @@
4040
)
4141
from google.protobuf import json_format
4242
from google.protobuf import struct_pb2
43+
import constants as test_constants
4344

44-
_TEST_BUCKET_NAME = "test-bucket"
45-
_TEST_GCS_PATH_WITHOUT_BUCKET = "path/to/folder"
46-
_TEST_GCS_PATH = f"{_TEST_BUCKET_NAME}/{_TEST_GCS_PATH_WITHOUT_BUCKET}"
47-
_TEST_GCS_PATH_WITH_TRAILING_SLASH = f"{_TEST_GCS_PATH}/"
45+
_TEST_BUCKET_NAME = test_constants.TrainingJobConstants._TEST_BUCKET_NAME
46+
_TEST_GCS_PATH_WITHOUT_BUCKET = (
47+
test_constants.TrainingJobConstants._TEST_GCS_PATH_WITHOUT_BUCKET
48+
)
49+
_TEST_GCS_PATH = test_constants.TrainingJobConstants._TEST_GCS_PATH
50+
_TEST_GCS_PATH_WITH_TRAILING_SLASH = (
51+
test_constants.TrainingJobConstants._TEST_GCS_PATH_WITH_TRAILING_SLASH
52+
)
4853
_TEST_PROJECT = "test-project"
4954

50-
_TEST_DATASET_DISPLAY_NAME = "test-dataset-display-name"
51-
_TEST_DATASET_NAME = "test-dataset-name"
52-
_TEST_DISPLAY_NAME = "test-display-name"
55+
_TEST_DATASET_DISPLAY_NAME = (
56+
test_constants.TrainingJobConstants._TEST_DATASET_DISPLAY_NAME
57+
)
58+
_TEST_DATASET_NAME = test_constants.TrainingJobConstants._TEST_DATASET_NAME
59+
_TEST_DISPLAY_NAME = test_constants.TrainingJobConstants._TEST_DISPLAY_NAME
5360
_TEST_METADATA_SCHEMA_URI_TABULAR = schema.dataset.metadata.tabular
5461
_TEST_METADATA_SCHEMA_URI_NONTABULAR = schema.dataset.metadata.image
5562

@@ -150,12 +157,12 @@
150157
struct_pb2.Value(),
151158
)
152159

153-
_TEST_DATASET_NAME = "test-dataset-name"
160+
_TEST_DATASET_NAME = test_constants.TrainingJobConstants._TEST_DATASET_NAME
154161

155-
_TEST_MODEL_DISPLAY_NAME = "model-display-name"
162+
_TEST_MODEL_DISPLAY_NAME = test_constants.TrainingJobConstants._TEST_MODEL_DISPLAY_NAME
156163

157-
_TEST_LABELS = {"key": "value"}
158-
_TEST_MODEL_LABELS = {"model_key": "model_value"}
164+
_TEST_LABELS = test_constants.ProjectConstants._TEST_LABELS
165+
_TEST_MODEL_LABELS = test_constants.TrainingJobConstants._TEST_MODEL_LABELS
159166

160167
_TEST_FRACTION_SPLIT_TRAINING = 0.6
161168
_TEST_FRACTION_SPLIT_VALIDATION = 0.2
@@ -169,11 +176,13 @@
169176
_TEST_MODEL_NAME = "projects/my-project/locations/us-central1/models/12345"
170177

171178
_TEST_PIPELINE_RESOURCE_NAME = (
172-
"projects/my-project/locations/us-central1/trainingPipelines/12345"
179+
test_constants.TrainingJobConstants._TEST_PIPELINE_RESOURCE_NAME
173180
)
174181

175182
# CMEK encryption
176-
_TEST_DEFAULT_ENCRYPTION_KEY_NAME = "key_default"
183+
_TEST_DEFAULT_ENCRYPTION_KEY_NAME = (
184+
test_constants.TrainingJobConstants._TEST_DEFAULT_ENCRYPTION_KEY_NAME
185+
)
177186
_TEST_DEFAULT_ENCRYPTION_SPEC = gca_encryption_spec.EncryptionSpec(
178187
kms_key_name=_TEST_DEFAULT_ENCRYPTION_KEY_NAME
179188
)

‎tests/unit/aiplatform/test_automl_text_training_jobs.py

Lines changed: 14 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -41,12 +41,15 @@
4141
from google.cloud.aiplatform.v1.schema.trainingjob import (
4242
definition_v1 as training_job_inputs,
4343
)
44+
import constants as test_constants
4445

4546
_TEST_PROJECT = "test-project"
4647
_TEST_LOCATION = "us-central1"
47-
_TEST_DATASET_DISPLAY_NAME = "test-dataset-display-name"
48-
_TEST_DATASET_NAME = "test-dataset-name"
49-
_TEST_DISPLAY_NAME = "test-display-name"
48+
_TEST_DATASET_DISPLAY_NAME = (
49+
test_constants.TrainingJobConstants._TEST_DATASET_DISPLAY_NAME
50+
)
51+
_TEST_DATASET_NAME = test_constants.TrainingJobConstants._TEST_DATASET_NAME
52+
_TEST_DISPLAY_NAME = test_constants.TrainingJobConstants._TEST_DISPLAY_NAME
5053
_TEST_METADATA_SCHEMA_URI_TEXT = schema.dataset.metadata.text
5154

5255
_TEST_PREDICTION_TYPE_CLASSIFICATION = "classification"
@@ -55,11 +58,11 @@
5558
_TEST_PREDICTION_TYPE_SENTIMENT = "sentiment"
5659
_TEST_SENTIMENT_MAX = 10
5760

58-
_TEST_DATASET_NAME = "test-dataset-name"
59-
_TEST_MODEL_DISPLAY_NAME = "model-display-name"
61+
_TEST_DATASET_NAME = test_constants.TrainingJobConstants._TEST_DATASET_NAME
62+
_TEST_MODEL_DISPLAY_NAME = test_constants.TrainingJobConstants._TEST_MODEL_DISPLAY_NAME
6063

61-
_TEST_LABELS = {"key": "value"}
62-
_TEST_MODEL_LABELS = {"model_key": "model_value"}
64+
_TEST_LABELS = test_constants.ProjectConstants._TEST_LABELS
65+
_TEST_MODEL_LABELS = test_constants.TrainingJobConstants._TEST_MODEL_LABELS
6366

6467
_TEST_MODEL_ID = "98777645321"
6568

@@ -86,11 +89,13 @@
8689
)
8790

8891
_TEST_PIPELINE_RESOURCE_NAME = (
89-
f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/trainingPipelines/12345"
92+
test_constants.TrainingJobConstants._TEST_PIPELINE_RESOURCE_NAME
9093
)
9194

9295
# CMEK encryption
93-
_TEST_DEFAULT_ENCRYPTION_KEY_NAME = "key_default"
96+
_TEST_DEFAULT_ENCRYPTION_KEY_NAME = (
97+
test_constants.TrainingJobConstants._TEST_DEFAULT_ENCRYPTION_KEY_NAME
98+
)
9499
_TEST_DEFAULT_ENCRYPTION_SPEC = gca_encryption_spec.EncryptionSpec(
95100
kms_key_name=_TEST_DEFAULT_ENCRYPTION_KEY_NAME
96101
)

‎tests/unit/aiplatform/test_automl_video_training_jobs.py

Lines changed: 14 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -41,12 +41,15 @@
4141
pipeline_state as gca_pipeline_state,
4242
training_pipeline as gca_training_pipeline,
4343
)
44+
import constants as test_constants
4445

4546
_TEST_PROJECT = "test-project"
4647
_TEST_LOCATION = "us-central1"
47-
_TEST_DATASET_DISPLAY_NAME = "test-dataset-display-name"
48-
_TEST_DATASET_NAME = "test-dataset-name"
49-
_TEST_DISPLAY_NAME = "test-display-name"
48+
_TEST_DATASET_DISPLAY_NAME = (
49+
test_constants.TrainingJobConstants._TEST_DATASET_DISPLAY_NAME
50+
)
51+
_TEST_DATASET_NAME = test_constants.TrainingJobConstants._TEST_DATASET_NAME
52+
_TEST_DISPLAY_NAME = test_constants.TrainingJobConstants._TEST_DISPLAY_NAME
5053
_TEST_METADATA_SCHEMA_URI_VIDEO = schema.dataset.metadata.video
5154

5255
_TEST_MODEL_TYPE_CLOUD = "CLOUD"
@@ -56,11 +59,11 @@
5659
_TEST_PREDICTION_TYPE_VCN = "classification"
5760
_TEST_PREDICTION_TYPE_VOR = "object_tracking"
5861

59-
_TEST_DATASET_NAME = "test-dataset-name"
60-
_TEST_MODEL_DISPLAY_NAME = "model-display-name"
62+
_TEST_DATASET_NAME = test_constants.TrainingJobConstants._TEST_DATASET_NAME
63+
_TEST_MODEL_DISPLAY_NAME = test_constants.TrainingJobConstants._TEST_MODEL_DISPLAY_NAME
6164

62-
_TEST_LABELS = {"key": "value"}
63-
_TEST_MODEL_LABELS = {"model_key": "model_value"}
65+
_TEST_LABELS = test_constants.ProjectConstants._TEST_LABELS
66+
_TEST_MODEL_LABELS = test_constants.TrainingJobConstants._TEST_MODEL_LABELS
6467

6568
_TEST_MODEL_ID = "98777645321" # TODO
6669

@@ -83,11 +86,13 @@
8386
)
8487

8588
_TEST_PIPELINE_RESOURCE_NAME = (
86-
f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/trainingPipelines/12345"
89+
test_constants.TrainingJobConstants._TEST_PIPELINE_RESOURCE_NAME
8790
)
8891

8992
# CMEK encryption
90-
_TEST_DEFAULT_ENCRYPTION_KEY_NAME = "key_default"
93+
_TEST_DEFAULT_ENCRYPTION_KEY_NAME = (
94+
test_constants.TrainingJobConstants._TEST_DEFAULT_ENCRYPTION_KEY_NAME
95+
)
9196
_TEST_DEFAULT_ENCRYPTION_SPEC = gca_encryption_spec.EncryptionSpec(
9297
kms_key_name=_TEST_DEFAULT_ENCRYPTION_KEY_NAME
9398
)

0 commit comments

Comments
 (0)