chore: internal clean up

tests/system/aiplatform/e2e_base.py

@@ -200,6 +200,8 @@ def tear_down_resources(self, shared_state: Dict[str, Any]):
                         aiplatform.Endpoint,
                         aiplatform.Featurestore,
                         aiplatform.MatchingEngineIndexEndpoint,
+                        aiplatform.BatchPredictionJob,
+                        aiplatform.ModelDeploymentMonitoringJob,
                     ),
                 ):
                     # For endpoint, undeploy model then delete endpoint

tests/system/aiplatform/test_featurestore.py

@@ -76,7 +76,6 @@ def test_create_get_list_featurestore(self, shared_state):
             featurestore_name=featurestore.resource_name
         )
         assert featurestore.resource_name == get_featurestore.resource_name
-
         list_featurestores = aiplatform.Featurestore.list()
         assert get_featurestore.resource_name in [
             featurestore.resource_name for featurestore in list_featurestores

tests/system/vertex_ray/test_ray_data.py

@@ -77,11 +77,15 @@ def create_bigquery_script(version: str):
 }
 
 
+@pytest.mark.usefixtures(
+    "prepare_bigquery_dataset",
+    "delete_bigquery_dataset",
+)
 class TestRayData(e2e_base.TestEndToEnd):
     _temp_prefix = "temp-ray-data"
 
     @pytest.mark.parametrize("cluster_ray_version", ["2.33", "2.42", "2.47"])
-    def test_ray_data(self, cluster_ray_version):
+    def test_ray_data(self, cluster_ray_version, shared_state):
         head_node_type = vertex_ray.Resources()
         worker_node_types = [
             vertex_ray.Resources(),
@@ -147,6 +151,14 @@ def test_ray_data(self, cluster_ray_version):
                     print(client.get_job_info(job_id).message)
                     raise RuntimeError("The Ray Job encountered an error and failed")
 
+        # Delete the bigquery dataset
+        version_suffix = cluster_ray_version.replace(".", "")
+        dataset_id = f"bugbashbq1.system_test_ray{version_suffix}_write"
+        bigquery_client = shared_state["bigquery_client"]
+        bigquery_client.delete_dataset(
+            dataset_id, delete_contents=True, not_found_ok=True
+        )
+
         vertex_ray.delete_ray_cluster(cluster_resource_name)
         # Ensure cluster was deleted
         for cluster in vertex_ray.list_ray_clusters():

tests/system/vertexai/test_batch_prediction.py

@@ -91,6 +91,10 @@ def test_batch_prediction_with_bq_input(self, shared_state):
             output_uri_prefix=f"bq://{shared_state['bigquery_dataset_id']}",
         )
         shared_state["resources"].append(job)
+        bigquery_client = shared_state["bigquery_client"]
+        bigquery_client.delete_dataset(
+            shared_state["bigquery_dataset_id"], delete_contents=True, not_found_ok=True
+        )
 
         assert (
             job.model_name == _GEMINI_MODEL_RESOURCE_NAME