diff --git a/tests/system/aiplatform/test_dataset.py b/tests/system/aiplatform/test_dataset.py index f4033bfb7e..f152ed0e32 100644 --- a/tests/system/aiplatform/test_dataset.py +++ b/tests/system/aiplatform/test_dataset.py @@ -131,7 +131,7 @@ def setup_method(self): @pytest.fixture() def storage_client(self): - yield storage.Client(project=e2e_base._PROJECT) + yield storage.Client(project=_TEST_PROJECT) @pytest.fixture() def staging_bucket(self, storage_client): @@ -174,7 +174,7 @@ def test_get_new_dataset_and_import(self, dataset_gapic_client): try: text_dataset = aiplatform.TextDataset.create( - display_name=f"temp_sdk_integration_test_create_text_dataset_{uuid.uuid4()}", + display_name=self._make_display_name(key="get_new_dataset_and_import"), ) my_dataset = aiplatform.TextDataset(dataset_name=text_dataset.name) @@ -189,7 +189,6 @@ def test_get_new_dataset_and_import(self, dataset_gapic_client): my_dataset.import_data( gcs_source=_TEST_TEXT_ENTITY_EXTRACTION_GCS_SOURCE, import_schema_uri=_TEST_TEXT_ENTITY_IMPORT_SCHEMA, - import_request_timeout=600.0, ) data_items_post_import = dataset_gapic_client.list_data_items( @@ -198,8 +197,7 @@ def test_get_new_dataset_and_import(self, dataset_gapic_client): assert len(list(data_items_post_import)) == 469 finally: - if text_dataset is not None: - text_dataset.delete() + text_dataset.delete() @vpcsc_config.skip_if_inside_vpcsc def test_create_and_import_image_dataset(self, dataset_gapic_client): @@ -208,7 +206,9 @@ def test_create_and_import_image_dataset(self, dataset_gapic_client): try: img_dataset = aiplatform.ImageDataset.create( - display_name=f"temp_sdk_integration_create_and_import_dataset_{uuid.uuid4()}", + display_name=self._make_display_name( + key="create_and_import_image_dataset" + ), gcs_source=_TEST_IMAGE_OBJECT_DETECTION_GCS_SOURCE, import_schema_uri=_TEST_IMAGE_OBJ_DET_IMPORT_SCHEMA, create_request_timeout=None, @@ -230,7 +230,7 @@ def test_create_tabular_dataset(self): try: tabular_dataset = aiplatform.TabularDataset.create( - display_name=f"temp_sdk_integration_create_and_import_dataset_{uuid.uuid4()}", + display_name=self._make_display_name(key="create_tabular_dataset"), gcs_source=[_TEST_TABULAR_CLASSIFICATION_GCS_SOURCE], create_request_timeout=None, ) @@ -250,13 +250,15 @@ def test_create_tabular_dataset(self): tabular_dataset.delete() def test_create_tabular_dataset_from_dataframe(self, bigquery_dataset): - bq_staging_table = f"bq://{e2e_base._PROJECT}.{bigquery_dataset.dataset_id}.test_table{uuid.uuid4()}" + bq_staging_table = f"bq://{_TEST_PROJECT}.{bigquery_dataset.dataset_id}.test_table{uuid.uuid4()}" try: tabular_dataset = aiplatform.TabularDataset.create_from_dataframe( df_source=_TEST_DATAFRAME, staging_path=bq_staging_table, - display_name=f"temp_sdk_integration_create_and_import_dataset_from_dataframe{uuid.uuid4()}", + display_name=self._make_display_name( + key="create_and_import_dataset_from_dataframe" + ), ) """Use the Dataset.create_from_dataframe() method to create a new tabular dataset. @@ -281,12 +283,14 @@ def test_create_tabular_dataset_from_dataframe_with_provided_schema( created and references the BQ source.""" try: - bq_staging_table = f"bq://{e2e_base._PROJECT}.{bigquery_dataset.dataset_id}.test_table{uuid.uuid4()}" + bq_staging_table = f"bq://{_TEST_PROJECT}.{bigquery_dataset.dataset_id}.test_table{uuid.uuid4()}" tabular_dataset = aiplatform.TabularDataset.create_from_dataframe( df_source=_TEST_DATAFRAME, staging_path=bq_staging_table, - display_name=f"temp_sdk_integration_create_and_import_dataset_from_dataframe{uuid.uuid4()}", + display_name=self._make_display_name( + key="create_and_import_dataset_from_dataframe" + ), bq_schema=_TEST_DATAFRAME_BQ_SCHEMA, )