Update cleanup logic in tests to empty the bucket instead of deleting the bucket

Mahalaxmibejugam · web-flow · commit 51122764f34c · 2025-12-01T12:32:49.000+05:30
diff --git a/docs/source/developer.rst b/docs/source/developer.rst
@@ -19,9 +19,9 @@ real GCS. A small number of tests run differently or are skipped.
 If you want to actually test against real GCS, then you should set
 STORAGE_EMULATOR_HOST to "https://storage.googleapis.com" and also
 provide appropriate GCSFS_TEST_BUCKET, GCSFS_TEST_VERSIONED_BUCKET
-(To use for tests that target GCS object versioning, this bucket must have object versioning enabled)
-and GCSFS_TEST_PROJECT, as well as setting your default google
-credentials (or providing them via the fsspec config).
+(To use for tests that target GCS object versioning, this bucket must have versioning enabled),
+GCSFS_ZONAL_TEST_BUCKET(To use for testing Rapid storage features) and GCSFS_TEST_PROJECT,
+as well as setting your default google credentials (or providing them via the fsspec config).
 
 .. _fake-gcs-server: https://github.com/fsouza/fake-gcs-server
 
diff --git a/gcsfs/tests/conftest.py b/gcsfs/tests/conftest.py
@@ -12,7 +12,7 @@
 from google.cloud import storage
 
 from gcsfs import GCSFileSystem
-from gcsfs.tests.settings import TEST_BUCKET, TEST_VERSIONED_BUCKET
+from gcsfs.tests.settings import TEST_BUCKET, TEST_VERSIONED_BUCKET, TEST_ZONAL_BUCKET
 
 files = {
     "test/accounts.1.json": (
@@ -60,7 +60,7 @@ def stop_docker(container):
         subprocess.call(["docker", "rm", "-f", "-v", cid])
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture(scope="session")
 def docker_gcs():
     if "STORAGE_EMULATOR_HOST" in os.environ:
         # assume using real API or otherwise have a server already set up
@@ -91,7 +91,7 @@ def docker_gcs():
     stop_docker(container)
 
 
-@pytest.fixture
+@pytest.fixture(scope="session")
 def gcs_factory(docker_gcs):
     params["endpoint_url"] = docker_gcs
 
@@ -102,44 +102,83 @@ def factory(**kwargs):
     return factory
 
 
+@pytest.fixture(scope="session")
+def buckets_to_delete():
+    """A set to keep track of buckets created during the test session."""
+    return set()
+
+
 @pytest.fixture
-def gcs(gcs_factory, populate=True):
+def gcs(gcs_factory, buckets_to_delete, populate=True):
     gcs = gcs_factory()
-    try:
-        # ensure we're empty.
-        try:
-            gcs.rm(TEST_BUCKET, recursive=True)
-        except FileNotFoundError:
-            pass
-        try:
+    try:  # ensure we're empty.
+        # Create the bucket if it doesn't exist, otherwise clean it.
+        if not gcs.exists(TEST_BUCKET):
             gcs.mkdir(TEST_BUCKET)
-        except Exception:
-            pass
+            buckets_to_delete.add(TEST_BUCKET)
+        else:
+            try:
+                gcs.rm(gcs.find(TEST_BUCKET))
+            except Exception as e:
+                logging.warning(f"Failed to empty bucket {TEST_BUCKET}: {e}")
 
         if populate:
             gcs.pipe({TEST_BUCKET + "/" + k: v for k, v in allfiles.items()})
         gcs.invalidate_cache()
         yield gcs
     finally:
-        try:
-            gcs.rm(gcs.find(TEST_BUCKET))
-            gcs.rm(TEST_BUCKET)
-        except:  # noqa: E722
-            pass
+        _cleanup_gcs(gcs)
 
 
-def _cleanup_gcs(gcs, is_real_gcs):
-    """Only remove the bucket/contents if we are NOT using the real GCS, logging a warning on failure."""
-    if is_real_gcs:
-        return
+def _cleanup_gcs(gcs):
+    """Clean the bucket contents, logging a warning on failure."""
     try:
-        gcs.rm(TEST_BUCKET, recursive=True)
+        gcs.rm(gcs.find(TEST_BUCKET))
     except Exception as e:
         logging.warning(f"Failed to clean up GCS bucket {TEST_BUCKET}: {e}")
 
 
+@pytest.fixture(scope="session", autouse=True)
+def final_cleanup(gcs_factory, buckets_to_delete):
+    """A session-scoped fixture to delete the test buckets after all tests are run."""
+    yield
+    # This code runs after the entire test session finishes
+    use_extended_gcs = os.getenv(
+        "GCSFS_EXPERIMENTAL_ZB_HNS_SUPPORT", "false"
+    ).lower() in (
+        "true",
+        "1",
+    )
+
+    if use_extended_gcs:
+        is_real_gcs = (
+            os.environ.get("STORAGE_EMULATOR_HOST") == "https://storage.googleapis.com"
+        )
+        mock_authentication_manager = (
+            patch("google.auth.default", return_value=(None, "fake-project"))
+            if not is_real_gcs
+            else nullcontext()
+        )
+    else:
+        mock_authentication_manager = nullcontext()
+
+    with mock_authentication_manager:
+        gcs = gcs_factory()
+        for bucket in buckets_to_delete:
+            # For real GCS, only delete if created by the test suite.
+            # For emulators, always delete.
+            try:
+                if gcs.exists(bucket):
+                    gcs.rm(bucket, recursive=True)
+                    logging.info(f"Cleaned up bucket: {bucket}")
+            except Exception as e:
+                logging.warning(
+                    f"Failed to perform final cleanup for bucket {bucket}: {e}"
+                )
+
+
 @pytest.fixture
-def extended_gcsfs(gcs_factory, populate=True):
+def extended_gcsfs(gcs_factory, buckets_to_delete, populate=True):
     # Check if we are running against a real GCS endpoint
     is_real_gcs = (
         os.environ.get("STORAGE_EMULATOR_HOST") == "https://storage.googleapis.com"
@@ -159,52 +198,62 @@ def extended_gcsfs(gcs_factory, populate=True):
             # Only create/delete/populate the bucket if we are NOT using the real GCS endpoint
             if not is_real_gcs:
                 try:
-                    extended_gcsfs.rm(TEST_BUCKET, recursive=True)
+                    extended_gcsfs.rm(TEST_ZONAL_BUCKET, recursive=True)
                 except FileNotFoundError:
                     pass
-                extended_gcsfs.mkdir(TEST_BUCKET)
+                extended_gcsfs.mkdir(TEST_ZONAL_BUCKET)
+                buckets_to_delete.add(TEST_ZONAL_BUCKET)
                 if populate:
                     extended_gcsfs.pipe(
-                        {TEST_BUCKET + "/" + k: v for k, v in allfiles.items()}
+                        {TEST_ZONAL_BUCKET + "/" + k: v for k, v in allfiles.items()}
                     )
             extended_gcsfs.invalidate_cache()
             yield extended_gcsfs
         finally:
-            _cleanup_gcs(extended_gcsfs, is_real_gcs)
+            _cleanup_gcs(extended_gcsfs)
 
 
 @pytest.fixture
-def gcs_versioned(gcs_factory):
+def gcs_versioned(gcs_factory, buckets_to_delete):
     gcs = gcs_factory()
     gcs.version_aware = True
     is_real_gcs = (
         os.environ.get("STORAGE_EMULATOR_HOST") == "https://storage.googleapis.com"
     )
     try:  # ensure we're empty.
+        # The versioned bucket might be created by `is_versioning_enabled`
+        # in test_core_versioned.py. We must register it for cleanup only if
+        # it was created by this test run.
+        try:
+            from gcsfs.tests.test_core_versioned import (
+                _VERSIONED_BUCKET_CREATED_BY_TESTS,
+            )
+
+            if _VERSIONED_BUCKET_CREATED_BY_TESTS:
+                buckets_to_delete.add(TEST_VERSIONED_BUCKET)
+        except ImportError:
+            pass  # test_core_versioned is not being run
         if is_real_gcs:
-            # For real GCS, we assume the bucket exists and only clean its contents.
-            try:
-                cleanup_versioned_bucket(gcs, TEST_VERSIONED_BUCKET)
-            except Exception as e:
-                logging.warning(
-                    f"Failed to empty versioned bucket {TEST_VERSIONED_BUCKET}: {e}"
-                )
+            cleanup_versioned_bucket(gcs, TEST_VERSIONED_BUCKET)
         else:
-            # For emulators, we delete and recreate the bucket for a clean state.
+            # For emulators, we delete and recreate the bucket for a clean state
             try:
                 gcs.rm(TEST_VERSIONED_BUCKET, recursive=True)
             except FileNotFoundError:
                 pass
             gcs.mkdir(TEST_VERSIONED_BUCKET, enable_versioning=True)
+            buckets_to_delete.add(TEST_VERSIONED_BUCKET)
         gcs.invalidate_cache()
         yield gcs
     finally:
+        # Ensure the bucket is empty after the test.
         try:
-            if not is_real_gcs:
-                gcs.rm(gcs.find(TEST_VERSIONED_BUCKET, versions=True))
-                gcs.rm(TEST_VERSIONED_BUCKET)
-        except:  # noqa: E722
-            pass
+            if is_real_gcs:
+                cleanup_versioned_bucket(gcs, TEST_VERSIONED_BUCKET)
+        except Exception as e:
+            logging.warning(
+                f"Failed to clean up versioned bucket {TEST_VERSIONED_BUCKET} after test: {e}"
+            )
 
 
 def cleanup_versioned_bucket(gcs, bucket_name, prefix=None):
diff --git a/gcsfs/tests/derived/gcsfs_fixtures.py b/gcsfs/tests/derived/gcsfs_fixtures.py
@@ -1,37 +1,35 @@
+import logging
+
 import fsspec
 import pytest
 from fsspec.tests.abstract import AbstractFixtures
 
 from gcsfs.core import GCSFileSystem
-from gcsfs.tests.conftest import allfiles
+from gcsfs.tests.conftest import _cleanup_gcs, allfiles
 from gcsfs.tests.settings import TEST_BUCKET
 
 
 class GcsfsFixtures(AbstractFixtures):
     @pytest.fixture(scope="class")
-    def fs(self, docker_gcs):
+    def fs(self, docker_gcs, buckets_to_delete):
         GCSFileSystem.clear_instance_cache()
         gcs = fsspec.filesystem("gcs", endpoint_url=docker_gcs)
-        try:
-            # ensure we're empty.
-            try:
-                gcs.rm(TEST_BUCKET, recursive=True)
-            except FileNotFoundError:
-                pass
-            try:
+        try:  # ensure we're empty.
+            # Create the bucket if it doesn't exist, otherwise clean it.
+            if not gcs.exists(TEST_BUCKET):
+                buckets_to_delete.add(TEST_BUCKET)
                 gcs.mkdir(TEST_BUCKET)
-            except Exception:
-                pass
+            else:
+                try:
+                    gcs.rm(gcs.find(TEST_BUCKET))
+                except Exception as e:
+                    logging.warning(f"Failed to empty bucket {TEST_BUCKET}: {e}")
 
             gcs.pipe({TEST_BUCKET + "/" + k: v for k, v in allfiles.items()})
             gcs.invalidate_cache()
             yield gcs
         finally:
-            try:
-                gcs.rm(gcs.find(TEST_BUCKET))
-                gcs.rm(TEST_BUCKET)
-            except:  # noqa: E722
-                pass
+            _cleanup_gcs(gcs)
 
     @pytest.fixture
     def fs_path(self):
diff --git a/gcsfs/tests/settings.py b/gcsfs/tests/settings.py
@@ -2,6 +2,7 @@
 
 TEST_BUCKET = os.getenv("GCSFS_TEST_BUCKET", "gcsfs_test")
 TEST_VERSIONED_BUCKET = os.getenv("GCSFS_TEST_VERSIONED_BUCKET", "gcsfs_test_versioned")
+TEST_ZONAL_BUCKET = os.getenv("GCSFS_ZONAL_TEST_BUCKET", "gcsfs_zonal_test")
 TEST_PROJECT = os.getenv("GCSFS_TEST_PROJECT", "project")
 TEST_REQUESTER_PAYS_BUCKET = f"{TEST_BUCKET}_req_pay"
 TEST_KMS_KEY = os.getenv(
diff --git a/gcsfs/tests/test_core_versioned.py b/gcsfs/tests/test_core_versioned.py
@@ -1,3 +1,4 @@
+import logging
 import os
 import posixpath
 
@@ -10,17 +11,28 @@
 a = TEST_VERSIONED_BUCKET + "/tmp/test/a"
 b = TEST_VERSIONED_BUCKET + "/tmp/test/b"
 
+# Flag to track if the bucket was created by this test run.
+_VERSIONED_BUCKET_CREATED_BY_TESTS = False
+
 
 def is_versioning_enabled():
     """
     Helper function to check if the test bucket has versioning enabled.
     Returns a tuple of (bool, reason_string).
     """
     # Don't skip when using an emulator, as we create the versioned bucket ourselves.
+    global _VERSIONED_BUCKET_CREATED_BY_TESTS
     if os.environ.get("STORAGE_EMULATOR_HOST") != "https://storage.googleapis.com":
         return True, ""
     try:
         gcs = GCSFileSystem(project=os.getenv("GCSFS_TEST_PROJECT", "project"))
+        if not gcs.exists(TEST_VERSIONED_BUCKET):
+            logging.info(
+                f"Creating versioned bucket for tests: {TEST_VERSIONED_BUCKET}"
+            )
+            gcs.mkdir(TEST_VERSIONED_BUCKET, enable_versioning=True)
+            _VERSIONED_BUCKET_CREATED_BY_TESTS = True
+
         client = storage.Client(
             credentials=gcs.credentials.credentials, project=gcs.project
         )
diff --git a/gcsfs/tests/test_extended_gcsfs.py b/gcsfs/tests/test_extended_gcsfs.py
@@ -11,17 +11,21 @@
 from google.cloud.storage.exceptions import DataCorruption
 
 from gcsfs.extended_gcsfs import BucketType
-from gcsfs.tests.conftest import a, b, c, csv_files, files, text_files
-from gcsfs.tests.settings import TEST_BUCKET
+from gcsfs.tests.conftest import csv_files, files, text_files
+from gcsfs.tests.settings import TEST_ZONAL_BUCKET
 
 file = "test/accounts.1.json"
-file_path = f"{TEST_BUCKET}/{file}"
+file_path = f"{TEST_ZONAL_BUCKET}/{file}"
 json_data = files[file]
 lines = io.BytesIO(json_data).readlines()
 file_size = len(json_data)
 
 REQUIRED_ENV_VAR = "GCSFS_EXPERIMENTAL_ZB_HNS_SUPPORT"
 
+a = TEST_ZONAL_BUCKET + "/tmp/test/a"
+b = TEST_ZONAL_BUCKET + "/tmp/test/b"
+c = TEST_ZONAL_BUCKET + "/tmp/test/c"
+
 # If the condition is True, only then tests in this file are run.
 should_run = os.getenv(REQUIRED_ENV_VAR, "false").lower() in (
     "true",
@@ -131,7 +135,7 @@ def test_read_block_zb(extended_gcsfs, zonal_mocks, subtests):
                 assert result == expected_data
                 if mocks:
                     mocks["sync_lookup_bucket_type"].assert_called_once_with(
-                        TEST_BUCKET
+                        TEST_ZONAL_BUCKET
                     )
                     if expected_data:
                         mocks["downloader"].download_ranges.assert_called_with(
@@ -143,7 +147,7 @@ def test_read_block_zb(extended_gcsfs, zonal_mocks, subtests):
 
 def test_read_small_zb(extended_gcsfs, zonal_mocks):
     csv_file = "2014-01-01.csv"
-    csv_file_path = f"{TEST_BUCKET}/{csv_file}"
+    csv_file_path = f"{TEST_ZONAL_BUCKET}/{csv_file}"
     csv_data = csv_files[csv_file]
 
     with zonal_mocks(csv_data) as mocks:
@@ -160,7 +164,9 @@ def test_read_small_zb(extended_gcsfs, zonal_mocks):
             # cache drop
             assert len(f.cache.cache) < len(out)
             if mocks:
-                mocks["sync_lookup_bucket_type"].assert_called_once_with(TEST_BUCKET)
+                mocks["sync_lookup_bucket_type"].assert_called_once_with(
+                    TEST_ZONAL_BUCKET
+                )
 
 
 def test_readline_zb(extended_gcsfs, zonal_mocks):
@@ -169,7 +175,7 @@ def test_readline_zb(extended_gcsfs, zonal_mocks):
     )
     for k, data in all_items:
         with zonal_mocks(data):
-            with extended_gcsfs.open("/".join([TEST_BUCKET, k]), "rb") as f:
+            with extended_gcsfs.open("/".join([TEST_ZONAL_BUCKET, k]), "rb") as f:
                 result = f.readline()
                 expected = data.split(b"\n")[0] + (b"\n" if data.count(b"\n") else b"")
             assert result == expected
