chore: fix Dataset system tests#1209
Conversation
Example: Remove shared state from test_dataset.py
sasha-gitg
left a comment
There was a problem hiding this comment.
Left a few comments. Approving to unblock release, we can resolve the comments afterwards. Thanks!
| shared_state["dataset_name"] = new_dataset.name | ||
| yield | ||
| def storage_client(self): | ||
| yield storage.Client(project=e2e_base._PROJECT) |
There was a problem hiding this comment.
Preference to reference either _TEST_PROJECT(line 130) or e2e_base._PROJECT but not both.
| assert len(list(data_items_post_import)) == 469 | ||
| assert len(list(data_items_post_import)) == 469 | ||
| finally: | ||
| if text_dataset is not None: |
There was a problem hiding this comment.
This condition doesn't seem necessary. Preference to remove it. If this line is present or not present, referencing text_dataset would throw.
| 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, |
There was a problem hiding this comment.
What's the purpose of the adding the timeout?
| aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) | ||
| try: | ||
| text_dataset = aiplatform.TextDataset.create( | ||
| display_name=f"temp_sdk_integration_test_create_text_dataset_{uuid.uuid4()}", |
There was a problem hiding this comment.
Preference to use this shared method for display names: https://github.com/googleapis/python-aiplatform/blob/main/tests/system/aiplatform/e2e_base.py#L50
| new_dataset = create_lro.result() | ||
| shared_state["dataset_name"] = new_dataset.name | ||
| yield | ||
| def staging_bucket(self, storage_client): |
There was a problem hiding this comment.
Preference to eventually rely on e2e_base implementation for deduplicaton of code.
Updated Dataset system tests to remove
shared_state, fixes failing system tests that were added withcreate_from_dataframemethod.