Merge remote-tracking branch 'origin/main' into zb-write-copy

Author: suni72

.github/workflows/ci.yml

@@ -42,10 +42,7 @@ jobs:
           pytest -vv -s \
           --log-format="%(asctime)s %(levelname)s %(message)s" \
           --log-date-format="%H:%M:%S" \
-          gcsfs/ \
-          --ignore=gcsfs/tests/test_extended_gcsfs.py \
-          --ignore=gcsfs/tests/test_zonal_file.py
-
+          gcsfs/
       - name: Run Extended Tests
         run: |
           pip install "git+https://github.com/googleapis/python-storage@main#egg=google-cloud-storage" --force-reinstall --no-cache-dir --upgrade

docs/source/developer.rst

@@ -18,9 +18,10 @@ 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).
 
 .. _fake-gcs-server: https://github.com/fsouza/fake-gcs-server
 

gcsfs/tests/conftest.py

@@ -7,9 +7,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": (
@@ -57,7 +58,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
@@ -88,7 +89,7 @@ def docker_gcs():
     stop_docker(container)
 
 
-@pytest.fixture
+@pytest.fixture(scope="session")
 def gcs_factory(docker_gcs):
     params["endpoint_url"] = docker_gcs
 
@@ -99,44 +100,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"
@@ -147,39 +187,96 @@ 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
-    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:
+                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)
+            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",