Extract bucket preparation logic into a reusable helper function in conftest.py

Author: Yonghui-Lee

gcsfs/tests/conftest.py

@@ -222,30 +222,14 @@ def gcs(gcs_factory, buckets_to_delete, populate_bucket):
         # Create the bucket if it doesn't exist, otherwise clean it.
         if not gcs.exists(TEST_BUCKET):
             gcs.mkdir(TEST_BUCKET)
+            # 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)
-            existing_files = {}
-        else:
-            existing_files = (
-                gcs.find(TEST_BUCKET, detail=True) if populate_bucket else {}
-            )
-            base_files = {TEST_BUCKET + "/" + k for k in allfiles.keys()}
-            if populate_bucket:
-                _cleanup_gcs(
-                    gcs, bucket=TEST_BUCKET, bucket_populated=True, excludes=base_files
-                )
-            else:
-                _cleanup_gcs(gcs, bucket=TEST_BUCKET, bucket_populated=False)
 
-        if populate_bucket:
-            files_to_pipe = {}
-            for k, v in allfiles.items():
-                remote_path = f"{TEST_BUCKET}/{k}"
-                if remote_path not in existing_files or existing_files[remote_path][
-                    "size"
-                ] != len(v):
-                    files_to_pipe[remote_path] = v
-            if files_to_pipe:
-                gcs.pipe(files_to_pipe)
+        _prepare_bucket(gcs, TEST_BUCKET, populate_bucket)
 
         gcs.invalidate_cache()
         yield gcs
@@ -277,7 +261,13 @@ def factory(**kwargs):
     yield factory
 
     for fs in created_instances:
-        _cleanup_gcs(fs, bucket_populated=populate_bucket)
+        base_files = {TEST_ZONAL_BUCKET + "/" + k for k in allfiles.keys()}
+        _cleanup_gcs(
+            fs,
+            bucket=TEST_ZONAL_BUCKET,
+            bucket_populated=populate_bucket,
+            excludes=base_files,
+        )
 
 
 @pytest.fixture
@@ -288,7 +278,13 @@ def extended_gcsfs(gcs_factory, buckets_to_delete, populate_bucket):
     try:
         yield extended_gcsfs
     finally:
-        _cleanup_gcs(extended_gcsfs, bucket_populated=populate_bucket)
+        base_files = {TEST_ZONAL_BUCKET + "/" + k for k in allfiles.keys()}
+        _cleanup_gcs(
+            extended_gcsfs,
+            bucket=TEST_ZONAL_BUCKET,
+            bucket_populated=populate_bucket,
+            excludes=base_files,
+        )
 
 
 def _cleanup_gcs(gcs, bucket=TEST_BUCKET, bucket_populated=True, excludes=None):
@@ -311,6 +307,26 @@ def _cleanup_gcs(gcs, bucket=TEST_BUCKET, bucket_populated=True, excludes=None):
         logging.warning(f"Failed to clean up GCS bucket {bucket}: {e}")
 
 
+def _prepare_bucket(gcs, bucket, populate_bucket):
+    """Clean bucket (excluding base files) and populate with missing/modified base files."""
+    if populate_bucket:
+        base_files = {bucket + "/" + k for k in allfiles.keys()}
+        _cleanup_gcs(gcs, bucket=bucket, bucket_populated=True, excludes=base_files)
+
+        existing_files = gcs.find(bucket, detail=True) if populate_bucket else {}
+        files_to_pipe = {}
+        for k, v in allfiles.items():
+            remote_path = f"{bucket}/{k}"
+            if remote_path not in existing_files or existing_files[remote_path][
+                "size"
+            ] != len(v):
+                files_to_pipe[remote_path] = v
+        if files_to_pipe:
+            gcs.pipe(files_to_pipe, finalize_on_close=True)
+    else:
+        _cleanup_gcs(gcs, bucket=bucket, bucket_populated=False)
+
+
 @pytest.fixture(scope="session", autouse=True)
 def final_cleanup(gcs_factory, buckets_to_delete):
     """
@@ -434,20 +450,7 @@ def _create_extended_gcsfs(gcs_factory, buckets_to_delete, populate_bucket, **kw
             extended_gcsfs.mkdir(TEST_BUCKET)
             buckets_to_delete.add(TEST_BUCKET)
     try:
-        if populate_bucket:
-            # To avoid hitting object mutation limits, only pipe files if they
-            # don't exist or if their size has changed.
-            existing_files = extended_gcsfs.find(TEST_ZONAL_BUCKET, detail=True)
-            files_to_pipe = {}
-            for k, v in allfiles.items():
-                remote_path = f"{TEST_ZONAL_BUCKET}/{k}"
-                if remote_path not in existing_files or existing_files[remote_path][
-                    "size"
-                ] != len(v):
-                    files_to_pipe[remote_path] = v
-
-            if files_to_pipe:
-                extended_gcsfs.pipe(files_to_pipe, finalize_on_close=True)
+        _prepare_bucket(extended_gcsfs, TEST_ZONAL_BUCKET, populate_bucket)
     except Exception as e:
         logging.warning(f"Failed to populate Zonal bucket: {e}")
 

gcsfs/tests/test_extended_hns_gcsfs.py

@@ -259,9 +259,8 @@ def test_hns_empty_folder_rename_success(self, gcs_hns, gcs_hns_mocks):
         path1 = f"{TEST_HNS_BUCKET}/empty_old_dir"
         path2 = f"{TEST_HNS_BUCKET}/empty_new_dir"
 
-        gcsfs.mkdir(path1)
-
         with gcs_hns_mocks(BucketType.HIERARCHICAL, gcsfs) as mocks:
+            gcsfs.mkdir(path1)
             # Configure mocks for the sequence of calls
             mocks["info"].side_effect = [
                 {"type": "directory", "name": path1},  # _mv check
@@ -274,7 +273,8 @@ def test_hns_empty_folder_rename_success(self, gcs_hns, gcs_hns_mocks):
             assert not gcsfs.exists(path1)
             assert gcsfs.exists(path2)
 
-            mocks["async_lookup_bucket_type"].assert_called_once_with(TEST_HNS_BUCKET)
+            assert mocks["async_lookup_bucket_type"].call_count == 2
+            mocks["async_lookup_bucket_type"].assert_called_with(TEST_HNS_BUCKET)
             expected_info_calls = [
                 mock.call(path1),  # from _mv
                 mock.call(path1),  # from exists(path1)
@@ -1568,9 +1568,8 @@ def test_hns_rmdir_success(self, gcs_hns, gcs_hns_mocks):
         dir_name = "empty_dir"
         dir_path = f"{TEST_HNS_BUCKET}/{dir_name}"
 
-        gcsfs.mkdir(f"{dir_path}/placeholder", create_parents=True)
-
         with gcs_hns_mocks(BucketType.HIERARCHICAL, gcsfs) as mocks:
+            gcsfs.mkdir(f"{dir_path}/placeholder", create_parents=True)
             # Configure mocks
             # 1. exists(dir_path) before rmdir should succeed.
             # 2. exists(dir_path) after rmdir should fail.
@@ -1583,7 +1582,8 @@ def test_hns_rmdir_success(self, gcs_hns, gcs_hns_mocks):
             gcsfs.rmdir(dir_path)
 
             assert not gcsfs.exists(dir_path)
-            mocks["async_lookup_bucket_type"].assert_called_once_with(TEST_HNS_BUCKET)
+            assert mocks["async_lookup_bucket_type"].call_count == 2
+            mocks["async_lookup_bucket_type"].assert_called_with(TEST_HNS_BUCKET)
             expected_request = self._get_delete_folder_request(dir_path)
             mocks["control_client"].delete_folder.assert_called_once_with(
                 request=expected_request, retry=mock.ANY, timeout=mock.ANY