Fixes for actual GCS bucket tests (fsspec#714)

Co-authored-by: Mahalaxmibejugam <[email protected]>

Author: ankitaluthra1

docs/source/developer.rst

@@ -18,9 +18,21 @@ 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 and GCSFS_TEST_PROJECT, as well
-as setting your default google credentials (or providing them via the
-fsspec config).
+provide appropriate GCSFS_TEST_BUCKET, GCSFS_TEST_VERSIONED_BUCKET
+(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).
+
+When running tests against a real GCS endpoint, you have two options for test buckets:
+
+- **Provide existing buckets**: If you specify buckets that already exist, the
+  test suite will manage objects *within* them (creating, modifying, and deleting
+  objects as needed). The buckets themselves will **not** be deleted upon completion.
+  **Warning**: The test suite will clear the contents of the bucket at the beginning and end of the
+  test run, so be sure to use a bucket that does not contain important data.
+- **Let the tests create buckets**: If you specify bucket names that do not exist,
+  the test suite will create them for the test run and automatically delete them
+  during final cleanup.
 
 .. _fake-gcs-server: https://github.com/fsouza/fake-gcs-server
 

gcsfs/tests/conftest.py

@@ -9,9 +9,10 @@
 import fsspec
 import pytest
 import requests
+from google.cloud import storage
 
 from gcsfs import GCSFileSystem
-from gcsfs.tests.settings import TEST_BUCKET
+from gcsfs.tests.settings import TEST_BUCKET, TEST_VERSIONED_BUCKET, TEST_ZONAL_BUCKET
 
 files = {
     "test/accounts.1.json": (
@@ -59,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
@@ -90,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
 
@@ -101,44 +102,102 @@ def factory(**kwargs):
     return factory
 
 
+@pytest.fixture(scope="session")
+def buckets_to_delete():
+    """
+    Provides a session-scoped set to track the names of GCS buckets that are
+    created by the test suite.
+
+    When tests run, they may create new GCS buckets. If these buckets are not
+    deleted, they will persist after the test run, leading to resource leakage.
+    This set acts as a registry of buckets that the `final_cleanup` fixture
+    should remove at the end of the entire 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
+            # By adding the bucket name to this set, we are marking it for
+            # deletion at the end of the test session. This ensures that if
+            # the test suite creates the bucket, it will also be responsible
+            # for deleting it. If the bucket already existed, we assume it's
+            # managed externally and should not be deleted by the tests.
+            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, auto-use fixture that deletes all buckets registered
+    in the `buckets_to_delete` set after the entire test session is complete.
+    """
+    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 if not using a real GCS endpoint,
+        # since grpc client in extended_gcsfs does not work with anon access
+        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:
+            # The cleanup logic attempts to delete every bucket that was
+            # added to the set during the session. For real GCS, only delete if
+            # created by the test suite.
+            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"
@@ -158,39 +217,102 @@ 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)
+                # When running against the emulator, the zonal test bucket is
+                # always created and added to the set for guaranteed cleanup.
+                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
-    try:
+    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:
-            gcs.rm(gcs.find(TEST_BUCKET, versions=True))
-        except FileNotFoundError:
-            pass
+            from gcsfs.tests.test_core_versioned import (
+                _VERSIONED_BUCKET_CREATED_BY_TESTS,
+            )
 
-        try:
-            gcs.mkdir(TEST_BUCKET, enable_versioning=True)
-        except Exception:
-            pass
+            if _VERSIONED_BUCKET_CREATED_BY_TESTS:
+                # If the versioned bucket was created by the test suite, it's added
+                # here for cleanup.
+                buckets_to_delete.add(TEST_VERSIONED_BUCKET)
+        except ImportError:
+            pass  # test_core_versioned is not being run
+        if is_real_gcs:
+            cleanup_versioned_bucket(gcs, TEST_VERSIONED_BUCKET)
+        else:
+            # 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)
+            # When using the emulator, the versioned bucket is always recreated
+            # and added to the cleanup set.
+            buckets_to_delete.add(TEST_VERSIONED_BUCKET)
         gcs.invalidate_cache()
         yield gcs
     finally:
+        # Ensure the bucket is empty after the test.
         try:
-            gcs.rm(gcs.find(TEST_BUCKET, versions=True))
-            gcs.rm(TEST_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):
+    """
+    Deletes all object versions in a bucket using the google-cloud-storage client,
+    ensuring it uses the same credentials as the gcsfs instance.
+    """
+    # Define a retry policy for API calls to handle rate limiting.
+    # This can retry on 429 Too Many Requests errors, which can happen
+    # when deleting many object versions quickly.
+    from google.api_core.retry import Retry
+
+    retry_policy = Retry(
+        initial=1.0,  # Initial delay in seconds
+        maximum=30.0,  # Maximum delay in seconds
+        multiplier=1.2,  # Backoff factor
+    )
+
+    client = storage.Client(
+        credentials=gcs.credentials.credentials, project=gcs.project
+    )
+
+    # List all blobs, including old versions
+    blobs_to_delete = list(client.list_blobs(bucket_name, versions=True, prefix=prefix))
+
+    if not blobs_to_delete:
+        logging.info("No object versions to delete in %s.", bucket_name)
+        return
+
+    logging.info(
+        "Deleting %d object versions from %s.", len(blobs_to_delete), bucket_name
+    )
+    time.sleep(2)
+    for blob in blobs_to_delete:
+        blob.delete(retry=retry_policy)
+
+    logging.info("Successfully deleted %d object versions.", len(blobs_to_delete))

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

gcsfs/tests/settings.py

@@ -1,8 +1,10 @@
 import os
 
 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 = "gcsfs_test_req_pay"
+TEST_REQUESTER_PAYS_BUCKET = f"{TEST_BUCKET}_req_pay"
 TEST_KMS_KEY = os.getenv(
     "GCSFS_TEST_KMS_KEY",
     f"projects/{TEST_PROJECT}/locations/us/keyRings/gcsfs_test/cryptKeys/gcsfs_test_key",