From d3924ab19922bd6cc5af2a581bd9f384821f5fbe Mon Sep 17 00:00:00 2001
From: Amy Wu <wuamy@google.com>
Date: Fri, 6 Feb 2026 14:27:01 -0800
Subject: [PATCH] chore: internal clean up

PiperOrigin-RevId: 866615593
---
 tests/system/aiplatform/e2e_base.py            |  2 ++
 tests/system/aiplatform/test_featurestore.py   |  1 -
 tests/system/vertex_ray/test_ray_data.py       | 14 +++++++++++++-
 tests/system/vertexai/test_batch_prediction.py |  4 ++++
 4 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/tests/system/aiplatform/e2e_base.py b/tests/system/aiplatform/e2e_base.py
index 2b1aefefa9..62ca536159 100644
--- a/tests/system/aiplatform/e2e_base.py
+++ b/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
diff --git a/tests/system/aiplatform/test_featurestore.py b/tests/system/aiplatform/test_featurestore.py
index 8aa1bd6f2f..95342a1244 100644
--- a/tests/system/aiplatform/test_featurestore.py
+++ b/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
diff --git a/tests/system/vertex_ray/test_ray_data.py b/tests/system/vertex_ray/test_ray_data.py
index aa7cf0ce11..a8d369bc53 100644
--- a/tests/system/vertex_ray/test_ray_data.py
+++ b/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():
diff --git a/tests/system/vertexai/test_batch_prediction.py b/tests/system/vertexai/test_batch_prediction.py
index 5a6c175586..8e486a079e 100644
--- a/tests/system/vertexai/test_batch_prediction.py
+++ b/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