From f9c770bb1a43fa7fe390286d7535266d3964d067 Mon Sep 17 00:00:00 2001
From: Albert Villanova del Moral
 <8515462+albertvillanova@users.noreply.github.com>
Date: Thu, 13 Apr 2023 12:54:13 +0200
Subject: [PATCH] Fix CI mock filesystem fixtures (#5740)

* Test mockfs fixture

* Do not pass clobber True to fsspec

* Create _register_custom_filesystems function

* Fix mock fs fixtures

* Use del instead of clear

* Revert 'Create _register_custom_filesystems function'

* Use fsspec registry instead of register_implementation

* Remove extra blank line
---
 tests/fixtures/fsspec.py                 | 10 +++++-----
 tests/test_filesystem.py                 | 11 +++++++++++
 tests/test_streaming_download_manager.py |  7 +++----
 3 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/tests/fixtures/fsspec.py b/tests/fixtures/fsspec.py
index 32f23d4e723..08dfb8bf2c3 100644
--- a/tests/fixtures/fsspec.py
+++ b/tests/fixtures/fsspec.py
@@ -2,9 +2,9 @@
 from pathlib import Path
 from unittest.mock import patch
 
-import fsspec
 import pytest
 from fsspec.implementations.local import AbstractFileSystem, LocalFileSystem, stringify_path
+from fsspec.registry import _registry as _fsspec_registry
 
 
 class MockFileSystem(AbstractFileSystem):
@@ -92,11 +92,11 @@ def _strip_protocol(cls, path):
 
 @pytest.fixture
 def mock_fsspec():
-    original_registry = fsspec.registry.copy()
-    fsspec.register_implementation("mock", MockFileSystem, clobber=True)
-    fsspec.register_implementation("tmp", TmpDirFileSystem, clobber=True)
+    _fsspec_registry["mock"] = MockFileSystem
+    _fsspec_registry["tmp"] = TmpDirFileSystem
     yield
-    fsspec.registry = original_registry
+    del _fsspec_registry["mock"]
+    del _fsspec_registry["tmp"]
 
 
 @pytest.fixture
diff --git a/tests/test_filesystem.py b/tests/test_filesystem.py
index cc6a0ad79ef..50b7939d805 100644
--- a/tests/test_filesystem.py
+++ b/tests/test_filesystem.py
@@ -2,12 +2,23 @@
 
 import fsspec
 import pytest
+from fsspec.registry import _registry as _fsspec_registry
 
 from datasets.filesystems import COMPRESSION_FILESYSTEMS, HfFileSystem, extract_path_from_uri, is_remote_filesystem
 
 from .utils import require_lz4, require_zstandard
 
 
+def test_mockfs(mockfs):
+    assert "mock" in _fsspec_registry
+    assert "bz2" in _fsspec_registry
+
+
+def test_non_mockfs():
+    assert "mock" not in _fsspec_registry
+    assert "bz2" in _fsspec_registry
+
+
 def test_extract_path_from_uri():
     mock_bucket = "mock-s3-bucket"
     dataset_path = f"s3://{mock_bucket}"
diff --git a/tests/test_streaming_download_manager.py b/tests/test_streaming_download_manager.py
index c33e6c0bbbd..08f64020e10 100644
--- a/tests/test_streaming_download_manager.py
+++ b/tests/test_streaming_download_manager.py
@@ -3,8 +3,8 @@
 import re
 from pathlib import Path
 
-import fsspec
 import pytest
+from fsspec.registry import _registry as _fsspec_registry
 from fsspec.spec import AbstractBufferedFile, AbstractFileSystem
 
 from datasets.download.streaming_download_manager import (
@@ -130,10 +130,9 @@ def _open(
 
 @pytest.fixture
 def mock_fsspec():
-    original_registry = fsspec.registry.copy()
-    fsspec.register_implementation("mock", DummyTestFS, clobber=True)
+    _fsspec_registry["mock"] = DummyTestFS
     yield
-    fsspec.registry = original_registry
+    del _fsspec_registry["mock"]
 
 
 def _readd_double_slash_removed_by_path(path_as_posix: str) -> str: