Skip to content

Commit 243c6f6

Browse files
authored
chore: cleanup on dataset system tests (#1214)
1 parent 50bdb01 commit 243c6f6

File tree

1 file changed

+15
-11
lines changed

1 file changed

+15
-11
lines changed

tests/system/aiplatform/test_dataset.py

+15-11
Original file line numberDiff line numberDiff line change
@@ -131,7 +131,7 @@ def setup_method(self):
131131

132132
@pytest.fixture()
133133
def storage_client(self):
134-
yield storage.Client(project=e2e_base._PROJECT)
134+
yield storage.Client(project=_TEST_PROJECT)
135135

136136
@pytest.fixture()
137137
def staging_bucket(self, storage_client):
@@ -174,7 +174,7 @@ def test_get_new_dataset_and_import(self, dataset_gapic_client):
174174

175175
try:
176176
text_dataset = aiplatform.TextDataset.create(
177-
display_name=f"temp_sdk_integration_test_create_text_dataset_{uuid.uuid4()}",
177+
display_name=self._make_display_name(key="get_new_dataset_and_import"),
178178
)
179179

180180
my_dataset = aiplatform.TextDataset(dataset_name=text_dataset.name)
@@ -189,7 +189,6 @@ def test_get_new_dataset_and_import(self, dataset_gapic_client):
189189
my_dataset.import_data(
190190
gcs_source=_TEST_TEXT_ENTITY_EXTRACTION_GCS_SOURCE,
191191
import_schema_uri=_TEST_TEXT_ENTITY_IMPORT_SCHEMA,
192-
import_request_timeout=600.0,
193192
)
194193

195194
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):
198197

199198
assert len(list(data_items_post_import)) == 469
200199
finally:
201-
if text_dataset is not None:
202-
text_dataset.delete()
200+
text_dataset.delete()
203201

204202
@vpcsc_config.skip_if_inside_vpcsc
205203
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):
208206

209207
try:
210208
img_dataset = aiplatform.ImageDataset.create(
211-
display_name=f"temp_sdk_integration_create_and_import_dataset_{uuid.uuid4()}",
209+
display_name=self._make_display_name(
210+
key="create_and_import_image_dataset"
211+
),
212212
gcs_source=_TEST_IMAGE_OBJECT_DETECTION_GCS_SOURCE,
213213
import_schema_uri=_TEST_IMAGE_OBJ_DET_IMPORT_SCHEMA,
214214
create_request_timeout=None,
@@ -230,7 +230,7 @@ def test_create_tabular_dataset(self):
230230

231231
try:
232232
tabular_dataset = aiplatform.TabularDataset.create(
233-
display_name=f"temp_sdk_integration_create_and_import_dataset_{uuid.uuid4()}",
233+
display_name=self._make_display_name(key="create_tabular_dataset"),
234234
gcs_source=[_TEST_TABULAR_CLASSIFICATION_GCS_SOURCE],
235235
create_request_timeout=None,
236236
)
@@ -250,13 +250,15 @@ def test_create_tabular_dataset(self):
250250
tabular_dataset.delete()
251251

252252
def test_create_tabular_dataset_from_dataframe(self, bigquery_dataset):
253-
bq_staging_table = f"bq://{e2e_base._PROJECT}.{bigquery_dataset.dataset_id}.test_table{uuid.uuid4()}"
253+
bq_staging_table = f"bq://{_TEST_PROJECT}.{bigquery_dataset.dataset_id}.test_table{uuid.uuid4()}"
254254

255255
try:
256256
tabular_dataset = aiplatform.TabularDataset.create_from_dataframe(
257257
df_source=_TEST_DATAFRAME,
258258
staging_path=bq_staging_table,
259-
display_name=f"temp_sdk_integration_create_and_import_dataset_from_dataframe{uuid.uuid4()}",
259+
display_name=self._make_display_name(
260+
key="create_and_import_dataset_from_dataframe"
261+
),
260262
)
261263

262264
"""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(
281283
created and references the BQ source."""
282284

283285
try:
284-
bq_staging_table = f"bq://{e2e_base._PROJECT}.{bigquery_dataset.dataset_id}.test_table{uuid.uuid4()}"
286+
bq_staging_table = f"bq://{_TEST_PROJECT}.{bigquery_dataset.dataset_id}.test_table{uuid.uuid4()}"
285287

286288
tabular_dataset = aiplatform.TabularDataset.create_from_dataframe(
287289
df_source=_TEST_DATAFRAME,
288290
staging_path=bq_staging_table,
289-
display_name=f"temp_sdk_integration_create_and_import_dataset_from_dataframe{uuid.uuid4()}",
291+
display_name=self._make_display_name(
292+
key="create_and_import_dataset_from_dataframe"
293+
),
290294
bq_schema=_TEST_DATAFRAME_BQ_SCHEMA,
291295
)
292296

0 commit comments

Comments
 (0)