From 64053a07716bb8bae322315f588dd463d16c9c21 Mon Sep 17 00:00:00 2001
From: yanghua <yanghua1127@gmail.com>
Date: Thu, 14 Nov 2024 11:40:24 +0800
Subject: [PATCH] Add more test cases about fsspec integration

---
 tosfs/tests/test_fsspec_integration.py | 54 ++++++++++++++++++++++++++
 1 file changed, 54 insertions(+)

diff --git a/tosfs/tests/test_fsspec_integration.py b/tosfs/tests/test_fsspec_integration.py
index d9facd5..e4606e1 100644
--- a/tosfs/tests/test_fsspec_integration.py
+++ b/tosfs/tests/test_fsspec_integration.py
@@ -18,6 +18,9 @@
 from fsspec.registry import known_implementations
 from tos import EnvCredentialsProvider
 
+from tosfs import TosFileSystem
+from tosfs.utils import random_str
+
 
 def test_fssepc_register():
     fsspec.register_implementation("tos", "tosfs.TosFileSystem")
@@ -41,3 +44,54 @@ def test_set_known_implementations():
         credentials_provider=EnvCredentialsProvider(),
     )
     assert len(tosfs.ls("")) > 0
+
+
+def test_fsspec_open(bucket, temporary_workspace):
+    known_implementations["tos"] = {"class": "tosfs.core.TosFileSystem"}
+
+    file = f"{random_str()}.txt"
+    content = "Hello TOSFS."
+    with fsspec.open(
+        f"tos://{bucket}/{temporary_workspace}/{file}",
+        endpoint_url=os.environ.get("TOS_ENDPOINT"),
+        region=os.environ.get("TOS_REGION"),
+        credentials_provider=EnvCredentialsProvider(),
+        mode="w",
+    ) as f:
+        f.write(content)
+
+    with fsspec.open(
+        f"tos://{bucket}/{temporary_workspace}/{file}",
+        endpoint_url=os.environ.get("TOS_ENDPOINT"),
+        region=os.environ.get("TOS_REGION"),
+        credentials_provider=EnvCredentialsProvider(),
+        mode="r",
+    ) as f:
+        assert f.read() == content
+
+
+class MyTosFileSystem(TosFileSystem):
+    """A sub class of TosFileSystem."""
+
+    def __init__(self):
+        """Init MyTosFileSystem."""
+        super().__init__(
+            endpoint_url=os.environ.get("TOS_ENDPOINT"),
+            region=os.environ.get("TOS_REGION"),
+            credentials_provider=EnvCredentialsProvider(),
+        )
+
+
+def test_customized_class(bucket, temporary_workspace):
+    known_implementations["tos"] = {
+        "class": "tosfs.tests.test_fsspec_integration.MyTosFileSystem"
+    }
+
+    file = f"{random_str()}.txt"
+    content = "Hello TOSFS."
+
+    with fsspec.open(f"tos://{bucket}/{temporary_workspace}/{file}", mode="w") as f:
+        f.write(content)
+
+    with fsspec.open(f"tos://{bucket}/{temporary_workspace}/{file}", mode="r") as f:
+        assert f.read() == content