chore: separating tests for efficiency

PiperOrigin-RevId: 524403464

tests/unit/aiplatform/constants.py

@@ -83,6 +83,15 @@ class TrainingJobConstants:
     _TEST_REDUCTION_SERVER_CONTAINER_URI = (
         "us-docker.pkg.dev/vertex-ai-restricted/training/reductionserver:latest"
     )
+    _TEST_DATASET_DISPLAY_NAME = "test-dataset-display-name"
+    _TEST_DATASET_NAME = "test-dataset-name"
+    _TEST_DISPLAY_NAME = "test-display-name"
+    _TEST_BUCKET_NAME = "test-bucket"
+    _TEST_GCS_PATH_WITHOUT_BUCKET = "path/to/folder"
+    _TEST_GCS_PATH = f"{_TEST_BUCKET_NAME}/{_TEST_GCS_PATH_WITHOUT_BUCKET}"
+    _TEST_GCS_PATH_WITH_TRAILING_SLASH = f"{_TEST_GCS_PATH}/"
+    _TEST_MODEL_DISPLAY_NAME = "model-display-name"
+    _TEST_MODEL_LABELS = {"model_key": "model_value"}
     _TEST_STAGING_BUCKET = "gs://test-staging-bucket"
     _TEST_DISPLAY_NAME = "my_job_1234"
     _TEST_BASE_OUTPUT_DIR = f"{_TEST_STAGING_BUCKET}/{_TEST_DISPLAY_NAME}"
@@ -146,6 +155,11 @@ class TrainingJobConstants:
     _TEST_TEST_FRACTION_SPLIT = 0.2
     _TEST_BOOT_DISK_TYPE_DEFAULT = "pd-ssd"
     _TEST_BOOT_DISK_SIZE_GB_DEFAULT = 100
+    # # DUPLICATE: THIS NEEDS TO BE MOVED TO THE TRAINING JOB TEST THAT USES IT
+    _TEST_PIPELINE_RESOURCE_NAME = (
+        "projects/my-project/locations/us-central1/trainingPipelines/12345"
+    )
+    _TEST_DEFAULT_ENCRYPTION_KEY_NAME = "key_default"
 
 
 @dataclasses.dataclass(frozen=True)

tests/unit/aiplatform/test_automl_forecasting_training_jobs.py

@@ -40,17 +40,26 @@
 )
 from google.protobuf import json_format
 from google.protobuf import struct_pb2
+import constants as test_constants
 
-_TEST_BUCKET_NAME = "test-bucket"
-_TEST_GCS_PATH_WITHOUT_BUCKET = "path/to/folder"
-_TEST_GCS_PATH = f"{_TEST_BUCKET_NAME}/{_TEST_GCS_PATH_WITHOUT_BUCKET}"
-_TEST_GCS_PATH_WITH_TRAILING_SLASH = f"{_TEST_GCS_PATH}/"
-_TEST_PROJECT = "test-project"
+_TEST_BUCKET_NAME = test_constants.TrainingJobConstants._TEST_BUCKET_NAME
+_TEST_GCS_PATH_WITHOUT_BUCKET = (
+    test_constants.TrainingJobConstants._TEST_GCS_PATH_WITHOUT_BUCKET
+)
+_TEST_GCS_PATH = test_constants.TrainingJobConstants._TEST_GCS_PATH
+_TEST_GCS_PATH_WITH_TRAILING_SLASH = (
+    test_constants.TrainingJobConstants._TEST_GCS_PATH_WITH_TRAILING_SLASH
+)
+_TEST_PROJECT = test_constants.ProjectConstants._TEST_PROJECT
 
-_TEST_DATASET_DISPLAY_NAME = "test-dataset-display-name"
-_TEST_DATASET_NAME = "test-dataset-name"
-_TEST_DISPLAY_NAME = "test-display-name"
-_TEST_TRAINING_CONTAINER_IMAGE = "gcr.io/test-training/container:image"
+_TEST_DATASET_DISPLAY_NAME = (
+    test_constants.TrainingJobConstants._TEST_DATASET_DISPLAY_NAME
+)
+_TEST_DATASET_NAME = test_constants.TrainingJobConstants._TEST_DATASET_NAME
+_TEST_DISPLAY_NAME = test_constants.TrainingJobConstants._TEST_DISPLAY_NAME
+_TEST_TRAINING_CONTAINER_IMAGE = (
+    test_constants.TrainingJobConstants._TEST_TRAINING_CONTAINER_IMAGE
+)
 _TEST_METADATA_SCHEMA_URI_TIMESERIES = schema.dataset.metadata.time_series
 _TEST_METADATA_SCHEMA_URI_NONTIMESERIES = schema.dataset.metadata.image
 
@@ -155,23 +164,25 @@
     struct_pb2.Value(),
 )
 
-_TEST_DATASET_NAME = "test-dataset-name"
+_TEST_DATASET_NAME = test_constants.TrainingJobConstants._TEST_DATASET_NAME
 
-_TEST_MODEL_DISPLAY_NAME = "model-display-name"
+_TEST_MODEL_DISPLAY_NAME = test_constants.TrainingJobConstants._TEST_MODEL_DISPLAY_NAME
 
-_TEST_LABELS = {"key": "value"}
-_TEST_MODEL_LABELS = {"model_key": "model_value"}
+_TEST_LABELS = test_constants.ProjectConstants._TEST_LABELS
+_TEST_MODEL_LABELS = test_constants.TrainingJobConstants._TEST_MODEL_LABELS
 
 _TEST_PREDEFINED_SPLIT_COLUMN_NAME = "split"
 
 _TEST_MODEL_NAME = "projects/my-project/locations/us-central1/models/12345"
 
 _TEST_PIPELINE_RESOURCE_NAME = (
-    "projects/my-project/locations/us-central1/trainingPipelines/12345"
+    test_constants.TrainingJobConstants._TEST_PIPELINE_RESOURCE_NAME
 )
 
 # CMEK encryption
-_TEST_DEFAULT_ENCRYPTION_KEY_NAME = "key_default"
+_TEST_DEFAULT_ENCRYPTION_KEY_NAME = (
+    test_constants.TrainingJobConstants._TEST_DEFAULT_ENCRYPTION_KEY_NAME
+)
 _TEST_DEFAULT_ENCRYPTION_SPEC = gca_encryption_spec.EncryptionSpec(
     kms_key_name=_TEST_DEFAULT_ENCRYPTION_KEY_NAME
 )

tests/unit/aiplatform/test_automl_image_training_jobs.py

@@ -41,12 +41,15 @@
     pipeline_state as gca_pipeline_state,
     training_pipeline as gca_training_pipeline,
 )
+import constants as test_constants
 
-_TEST_PROJECT = "test-project"
-_TEST_LOCATION = "us-central1"
-_TEST_DATASET_DISPLAY_NAME = "test-dataset-display-name"
-_TEST_DATASET_NAME = "test-dataset-name"
-_TEST_DISPLAY_NAME = "test-display-name"
+_TEST_PROJECT = test_constants.ProjectConstants._TEST_PROJECT
+_TEST_LOCATION = test_constants.ProjectConstants._TEST_LOCATION
+_TEST_DATASET_DISPLAY_NAME = (
+    test_constants.TrainingJobConstants._TEST_DATASET_DISPLAY_NAME
+)
+_TEST_DATASET_NAME = test_constants.TrainingJobConstants._TEST_DATASET_NAME
+_TEST_DISPLAY_NAME = test_constants.TrainingJobConstants._TEST_DISPLAY_NAME
 _TEST_METADATA_SCHEMA_URI_IMAGE = schema.dataset.metadata.image
 
 _TEST_TRAINING_BUDGET_MILLI_NODE_HOURS = 7500
@@ -57,12 +60,11 @@
 _TEST_PREDICTION_TYPE_ICN = "classification"
 _TEST_PREDICTION_TYPE_IOD = "object_detection"
 
-_TEST_DATASET_NAME = "test-dataset-name"
-_TEST_MODEL_DISPLAY_NAME = "model-display-name"
+_TEST_MODEL_DISPLAY_NAME = test_constants.TrainingJobConstants._TEST_MODEL_DISPLAY_NAME
 _TEST_MODEL_ID = "98777645321"
 
-_TEST_LABELS = {"key": "value"}
-_TEST_MODEL_LABELS = {"model_key": "model_value"}
+_TEST_LABELS = test_constants.ProjectConstants._TEST_LABELS
+_TEST_MODEL_LABELS = test_constants.TrainingJobConstants._TEST_MODEL_LABELS
 
 _TEST_TRAINING_TASK_INPUTS = json_format.ParseDict(
     {
@@ -109,11 +111,13 @@
 )
 
 _TEST_PIPELINE_RESOURCE_NAME = (
-    f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/trainingPipelines/12345"
+    test_constants.TrainingJobConstants._TEST_PIPELINE_RESOURCE_NAME
 )
 
 # CMEK encryption
-_TEST_DEFAULT_ENCRYPTION_KEY_NAME = "key_default"
+_TEST_DEFAULT_ENCRYPTION_KEY_NAME = (
+    test_constants.TrainingJobConstants._TEST_DEFAULT_ENCRYPTION_KEY_NAME
+)
 _TEST_DEFAULT_ENCRYPTION_SPEC = gca_encryption_spec.EncryptionSpec(
     kms_key_name=_TEST_DEFAULT_ENCRYPTION_KEY_NAME
 )

tests/unit/aiplatform/test_automl_tabular_training_jobs.py

@@ -40,16 +40,23 @@
 )
 from google.protobuf import json_format
 from google.protobuf import struct_pb2
+import constants as test_constants
 
-_TEST_BUCKET_NAME = "test-bucket"
-_TEST_GCS_PATH_WITHOUT_BUCKET = "path/to/folder"
-_TEST_GCS_PATH = f"{_TEST_BUCKET_NAME}/{_TEST_GCS_PATH_WITHOUT_BUCKET}"
-_TEST_GCS_PATH_WITH_TRAILING_SLASH = f"{_TEST_GCS_PATH}/"
+_TEST_BUCKET_NAME = test_constants.TrainingJobConstants._TEST_BUCKET_NAME
+_TEST_GCS_PATH_WITHOUT_BUCKET = (
+    test_constants.TrainingJobConstants._TEST_GCS_PATH_WITHOUT_BUCKET
+)
+_TEST_GCS_PATH = test_constants.TrainingJobConstants._TEST_GCS_PATH
+_TEST_GCS_PATH_WITH_TRAILING_SLASH = (
+    test_constants.TrainingJobConstants._TEST_GCS_PATH_WITH_TRAILING_SLASH
+)
 _TEST_PROJECT = "test-project"
 
-_TEST_DATASET_DISPLAY_NAME = "test-dataset-display-name"
-_TEST_DATASET_NAME = "test-dataset-name"
-_TEST_DISPLAY_NAME = "test-display-name"
+_TEST_DATASET_DISPLAY_NAME = (
+    test_constants.TrainingJobConstants._TEST_DATASET_DISPLAY_NAME
+)
+_TEST_DATASET_NAME = test_constants.TrainingJobConstants._TEST_DATASET_NAME
+_TEST_DISPLAY_NAME = test_constants.TrainingJobConstants._TEST_DISPLAY_NAME
 _TEST_METADATA_SCHEMA_URI_TABULAR = schema.dataset.metadata.tabular
 _TEST_METADATA_SCHEMA_URI_NONTABULAR = schema.dataset.metadata.image
 
@@ -150,12 +157,12 @@
     struct_pb2.Value(),
 )
 
-_TEST_DATASET_NAME = "test-dataset-name"
+_TEST_DATASET_NAME = test_constants.TrainingJobConstants._TEST_DATASET_NAME
 
-_TEST_MODEL_DISPLAY_NAME = "model-display-name"
+_TEST_MODEL_DISPLAY_NAME = test_constants.TrainingJobConstants._TEST_MODEL_DISPLAY_NAME
 
-_TEST_LABELS = {"key": "value"}
-_TEST_MODEL_LABELS = {"model_key": "model_value"}
+_TEST_LABELS = test_constants.ProjectConstants._TEST_LABELS
+_TEST_MODEL_LABELS = test_constants.TrainingJobConstants._TEST_MODEL_LABELS
 
 _TEST_FRACTION_SPLIT_TRAINING = 0.6
 _TEST_FRACTION_SPLIT_VALIDATION = 0.2
@@ -169,11 +176,13 @@
 _TEST_MODEL_NAME = "projects/my-project/locations/us-central1/models/12345"
 
 _TEST_PIPELINE_RESOURCE_NAME = (
-    "projects/my-project/locations/us-central1/trainingPipelines/12345"
+    test_constants.TrainingJobConstants._TEST_PIPELINE_RESOURCE_NAME
 )
 
 # CMEK encryption
-_TEST_DEFAULT_ENCRYPTION_KEY_NAME = "key_default"
+_TEST_DEFAULT_ENCRYPTION_KEY_NAME = (
+    test_constants.TrainingJobConstants._TEST_DEFAULT_ENCRYPTION_KEY_NAME
+)
 _TEST_DEFAULT_ENCRYPTION_SPEC = gca_encryption_spec.EncryptionSpec(
     kms_key_name=_TEST_DEFAULT_ENCRYPTION_KEY_NAME
 )

tests/unit/aiplatform/test_automl_text_training_jobs.py

@@ -41,12 +41,15 @@
 from google.cloud.aiplatform.v1.schema.trainingjob import (
     definition_v1 as training_job_inputs,
 )
+import constants as test_constants
 
 _TEST_PROJECT = "test-project"
 _TEST_LOCATION = "us-central1"
-_TEST_DATASET_DISPLAY_NAME = "test-dataset-display-name"
-_TEST_DATASET_NAME = "test-dataset-name"
-_TEST_DISPLAY_NAME = "test-display-name"
+_TEST_DATASET_DISPLAY_NAME = (
+    test_constants.TrainingJobConstants._TEST_DATASET_DISPLAY_NAME
+)
+_TEST_DATASET_NAME = test_constants.TrainingJobConstants._TEST_DATASET_NAME
+_TEST_DISPLAY_NAME = test_constants.TrainingJobConstants._TEST_DISPLAY_NAME
 _TEST_METADATA_SCHEMA_URI_TEXT = schema.dataset.metadata.text
 
 _TEST_PREDICTION_TYPE_CLASSIFICATION = "classification"
@@ -55,11 +58,11 @@
 _TEST_PREDICTION_TYPE_SENTIMENT = "sentiment"
 _TEST_SENTIMENT_MAX = 10
 
-_TEST_DATASET_NAME = "test-dataset-name"
-_TEST_MODEL_DISPLAY_NAME = "model-display-name"
+_TEST_DATASET_NAME = test_constants.TrainingJobConstants._TEST_DATASET_NAME
+_TEST_MODEL_DISPLAY_NAME = test_constants.TrainingJobConstants._TEST_MODEL_DISPLAY_NAME
 
-_TEST_LABELS = {"key": "value"}
-_TEST_MODEL_LABELS = {"model_key": "model_value"}
+_TEST_LABELS = test_constants.ProjectConstants._TEST_LABELS
+_TEST_MODEL_LABELS = test_constants.TrainingJobConstants._TEST_MODEL_LABELS
 
 _TEST_MODEL_ID = "98777645321"
 
@@ -86,11 +89,13 @@
 )
 
 _TEST_PIPELINE_RESOURCE_NAME = (
-    f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/trainingPipelines/12345"
+    test_constants.TrainingJobConstants._TEST_PIPELINE_RESOURCE_NAME
 )
 
 # CMEK encryption
-_TEST_DEFAULT_ENCRYPTION_KEY_NAME = "key_default"
+_TEST_DEFAULT_ENCRYPTION_KEY_NAME = (
+    test_constants.TrainingJobConstants._TEST_DEFAULT_ENCRYPTION_KEY_NAME
+)
 _TEST_DEFAULT_ENCRYPTION_SPEC = gca_encryption_spec.EncryptionSpec(
     kms_key_name=_TEST_DEFAULT_ENCRYPTION_KEY_NAME
 )

tests/unit/aiplatform/test_automl_video_training_jobs.py

@@ -41,12 +41,15 @@
     pipeline_state as gca_pipeline_state,
     training_pipeline as gca_training_pipeline,
 )
+import constants as test_constants
 
 _TEST_PROJECT = "test-project"
 _TEST_LOCATION = "us-central1"
-_TEST_DATASET_DISPLAY_NAME = "test-dataset-display-name"
-_TEST_DATASET_NAME = "test-dataset-name"
-_TEST_DISPLAY_NAME = "test-display-name"
+_TEST_DATASET_DISPLAY_NAME = (
+    test_constants.TrainingJobConstants._TEST_DATASET_DISPLAY_NAME
+)
+_TEST_DATASET_NAME = test_constants.TrainingJobConstants._TEST_DATASET_NAME
+_TEST_DISPLAY_NAME = test_constants.TrainingJobConstants._TEST_DISPLAY_NAME
 _TEST_METADATA_SCHEMA_URI_VIDEO = schema.dataset.metadata.video
 
 _TEST_MODEL_TYPE_CLOUD = "CLOUD"
@@ -56,11 +59,11 @@
 _TEST_PREDICTION_TYPE_VCN = "classification"
 _TEST_PREDICTION_TYPE_VOR = "object_tracking"
 
-_TEST_DATASET_NAME = "test-dataset-name"
-_TEST_MODEL_DISPLAY_NAME = "model-display-name"
+_TEST_DATASET_NAME = test_constants.TrainingJobConstants._TEST_DATASET_NAME
+_TEST_MODEL_DISPLAY_NAME = test_constants.TrainingJobConstants._TEST_MODEL_DISPLAY_NAME
 
-_TEST_LABELS = {"key": "value"}
-_TEST_MODEL_LABELS = {"model_key": "model_value"}
+_TEST_LABELS = test_constants.ProjectConstants._TEST_LABELS
+_TEST_MODEL_LABELS = test_constants.TrainingJobConstants._TEST_MODEL_LABELS
 
 _TEST_MODEL_ID = "98777645321"  # TODO
 
@@ -83,11 +86,13 @@
 )
 
 _TEST_PIPELINE_RESOURCE_NAME = (
-    f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/trainingPipelines/12345"
+    test_constants.TrainingJobConstants._TEST_PIPELINE_RESOURCE_NAME
 )
 
 # CMEK encryption
-_TEST_DEFAULT_ENCRYPTION_KEY_NAME = "key_default"
+_TEST_DEFAULT_ENCRYPTION_KEY_NAME = (
+    test_constants.TrainingJobConstants._TEST_DEFAULT_ENCRYPTION_KEY_NAME
+)
 _TEST_DEFAULT_ENCRYPTION_SPEC = gca_encryption_spec.EncryptionSpec(
     kms_key_name=_TEST_DEFAULT_ENCRYPTION_KEY_NAME
 )