Merge pull request #72 from RRosio/transfer_tests

Update upload and download tests

Author: RRosio

conftest.py

@@ -5,6 +5,10 @@
 from moto.moto_server.threaded_moto_server import ThreadedMotoServer
 from jupyter_fsspec.file_manager import FileSystemManager
 
+PORT = 5555
+ENDPOINT_URI = "http://127.0.0.1:%s/" % PORT
+
+
 pytest_plugins = [
     "pytest_jupyter.jupyter_server",
     "jupyter_server.pytest_plugin",
@@ -16,28 +20,39 @@
 def setup_tmp_local(tmp_path: Path):
     local_root = tmp_path / "test"
     local_root.mkdir()
+    nested_dir = local_root / "nested"
+    nested_dir.mkdir()
+    nested_file1 = nested_dir / ".empty"
+    nested_file1.touch()
+    nested_file2 = nested_dir / ".keep"
+    nested_file2.touch()
     local_file = local_root / "file_loc.txt"
     local_file.touch()
+
     local_empty_root = tmp_path / "empty"
     local_empty_root.mkdir()
 
-    return [local_root, local_empty_root]
+    yield [local_root, local_empty_root]
 
 
 @pytest.fixture(scope="function", autouse=True)
 def setup_config_file_fs(tmp_path: Path, setup_tmp_local):
     tmp_local = setup_tmp_local[0]
+    items_tmp_local = list(tmp_local.iterdir())
+    print(f"items_tmp_local: {items_tmp_local}")
     empty_tmp_local = setup_tmp_local[1]
     config_dir = tmp_path / "config"
     config_dir.mkdir(exist_ok=True)
 
     yaml_content = f"""sources:
   - name: "TestSourceAWS"
     path: "s3://my-test-bucket/"
-    additional_options:
+    kwargs:
       anon: false
       key: "my-access-key"
       secret: "my-secret-key"
+      client_kwargs:
+        endpoint_url: "{ENDPOINT_URI}"
   - name: "TestDir"
     path: "{tmp_local}"
     protocol: "local"
@@ -57,11 +72,11 @@ def setup_config_file_fs(tmp_path: Path, setup_tmp_local):
         print(f"Patching jupyter_config_dir to: {config_dir}")
         fs_manager = FileSystemManager(config_file="jupyter-fsspec.yaml")
 
-    return fs_manager
+    yield fs_manager
 
 
 @pytest.fixture(scope="function")
-def fs_manager_instance(setup_config_file_fs):
+def fs_manager_instance(setup_config_file_fs, s3_client):
     fs_manager = setup_config_file_fs
     fs_info = fs_manager.get_filesystem_by_protocol("memory")
     mem_fs = fs_info["info"]["instance"]
@@ -88,8 +103,7 @@ def fs_manager_instance(setup_config_file_fs):
         print("In memory filesystem NOT FOUND")
 
     if mem_fs.exists(f"{mem_root_path}/test_dir/file1.txt"):
-        file_info = mem_fs.info(f"{mem_root_path}/test_dir/file1.txt")
-        print(f"File exists. size: {file_info}")
+        mem_fs.info(f"{mem_root_path}/test_dir/file1.txt")
     else:
         print("File does not exist!")
     return fs_manager
@@ -98,120 +112,45 @@ def fs_manager_instance(setup_config_file_fs):
 def get_boto3_client():
     from botocore.session import Session
 
-    # NB: we use the sync botocore client for setup
     session = Session()
+    return session.create_client(
+        "s3",
+        endpoint_url=ENDPOINT_URI,
+        aws_access_key_id="my-access-key",
+        aws_secret_access_key="my-secret-key",
+    )
 
-    endpoint_uri = "http://127.0.0.1:%s/" % "5555"
-    return session.create_client("s3", endpoint_url=endpoint_uri)
-
-
-@pytest.fixture(scope="function")
-def s3_client(mock_s3_fs):
-    s3_client = get_boto3_client()
-    s3_client.create_bucket(Bucket="my-test-bucket")
-    return s3_client
-
-
-@pytest.fixture(scope="function")
-def s3_fs_manager_instance(setup_config_file_fs):
-    fs_manager = setup_config_file_fs
-    # fs_info = fs_manager.get_filesystem_by_protocol("s3")
-    # key = fs_info["key"]
-    # fs = fs_info["info"]["instance"]
-    # root_path = fs_info["info"]["path"]
-
-    # endpoint_uri = "http://127.0.0.1:%s/" % "5555"
-    # fs = fsspec.filesystem('s3', asynchronous=True, anon=False, client_kwargs={'endpoint_url': endpoint_uri})
-    return fs_manager
-
-
-@pytest.fixture(params=["memory", "local", "s3"])
-def filesystem_protocol(request):
-    return request.param
 
-
-@pytest.fixture(scope="function")
-def populated_file_system(filesystem_protocol):
-    fs_manager = FileSystemManager(config_file="jupyter-fsspec.yaml")
-    fs_protocol = filesystem_protocol
-    fs_info = fs_manager.get_filesystem_by_protocol(fs_protocol)
-    fs = fs_info["info"]["instance"]
-
-    if fs:
-        # Delete any existting directories
-        # Populate the filesystem
-        # mkdir => root_path + 'rootA'
-        # mkdir => root_path + 'rootB'
-        # touch => root_path + 'file1.txt'
-        # touch => root_path + 'rootA' + 'file_in_rootA.txt'
-        print(f"valid filesystem: {fs}")
-    else:
-        print(f"invalid filesystem: {fs}")
-    return {"fs_protocol": fs_protocol, "fs_manager": fs_manager}
-
-
-# TODO: Update this fixture from s3fs
-@pytest.fixture(scope="function")
-def mock_s3_fs():
-    # This fixture is module-scoped, meaning that we can re-use the MotoServer across all tests
-    server = ThreadedMotoServer(ip_address="127.0.0.1", port=5555)
+@pytest.fixture(scope="module")
+def s3_base():
+    server = ThreadedMotoServer(ip_address="127.0.0.1", port=PORT)
     server.start()
     if "AWS_SECRET_ACCESS_KEY" not in os.environ:
-        os.environ["AWS_SECRET_ACCESS_KEY"] = "foo"
+        os.environ["AWS_SECRET_ACCESS_KEY"] = "my-accesss-key"
     if "AWS_ACCESS_KEY_ID" not in os.environ:
-        os.environ["AWS_ACCESS_KEY_ID"] = "foo"
-    # aws_session_token=os.environ["AWS_SESSION_TOKEN"]
-    if "AWS_SESSION_TOKEN" not in os.environ:
-        os.environ["AWS_SESSION_TOKEN"] = "foo"
+        os.environ["AWS_ACCESS_KEY_ID"] = "my-secret-key"
+
     print("server up")
     yield
     print("moto done")
     server.stop()
 
 
-@pytest.fixture(scope="function")
-def fs_manager_instance_parameterized(populated_file_system):
-    fs_ret = populated_file_system
-    fs_protocol = fs_ret["fs_protocol"]
-    fs_manager = fs_ret["fs_manager"]
-    fs_info = fs_manager.get_filesystem_by_protocol(fs_protocol)
-    fs = fs_info["info"]["instance"]
-    root_path = fs_info["info"]["path"]
-
-    # fs_info = fs_manager.get_filesystem_by_protocol('local')
-    # key = fs_info['key']
-    # fs = fs_info['info']['instance']
-    # local_root_path = fs_info['info']['path']
-
-    if fs:
-        # TODO: Update file creation FOR PATHS!!!
-        if fs.exists(f"{root_path}/test_dir"):
-            print(f"{root_path}/test_dir EXISTS!!!!")
-            # fs.rm(f'{root_path}/test_dir', recursive=True)
-        if fs.exists(f"{root_path}/second_dir"):
-            print(f"{root_path}/second_dir EXISTS!!!!")
-            # fs.rm('/my_dir/second_dir', recursive=True)
-
-        fs.touch(f"{root_path}/file_in_root.txt")
-        with fs.open(f"{root_path}/file_in_root.txt", "wb") as f:
-            f.write("Root file content".encode())
-
-        # fs.mkdir('/my_dir/test_dir', exist_ok=True)
-        # fs.mkdir('/my_dir/second_dir', exist_ok=True)
-        # # fs.mkdir('/my_dir/second_dir/subdir', exist_ok=True)
-        # fs.touch('/my_dir/test_dir/file1.txt')
-        # with fs.open('/my_dir/test_dir/file1.txt', "wb") as f:
-        #     f.write("Test content".encode())
-        #     f.close()
-    else:
-        print(f"Filesystem of protocol {fs_protocol} NOT FOUND")
-
-    if fs.exists(f"{root_path}test_dir/file1.txt"):
-        file_info = fs.info(f"/{root_path}/test_dir/file1.txt")
-        print(f"File exists. size: {file_info}")
-    else:
-        print("File does not exist!")
-    return fs_manager
+@pytest.fixture(scope="module")
+def s3_client(s3_base):
+    client = get_boto3_client()
+    client.create_bucket(Bucket="my-test-bucket", ACL="public-read")
+    client.put_object(
+        Body=b"Hello, World1!", Bucket="my-test-bucket", Key="bucket-filename1.txt"
+    )
+    client.put_object(
+        Body=b"Hello, World2!", Bucket="my-test-bucket", Key="some/bucket-filename2.txt"
+    )
+    client.put_object(
+        Body=b"Hello, World3!", Bucket="my-test-bucket", Key="some/bucket-filename3.txt"
+    )
+
+    yield client
 
 
 @pytest.fixture

jupyter_fsspec/file_manager.py

@@ -211,3 +211,8 @@ def get_filesystem_by_protocol(self, fs_protocol):
             if fs_info.get("protocol") == fs_protocol:
                 return {"key": encoded_key, "info": fs_info}
         return None
+
+    def get_filesystem_protocol(self, key):
+        filesystem_rep = self.filesystems.get(key)
+        print(f"filesystem_rep: {filesystem_rep}")
+        return filesystem_rep["protocol"] + "://"

jupyter_fsspec/handlers.py

@@ -1,11 +1,14 @@
+from .file_manager import FileSystemManager
 from jupyter_server.base.handlers import APIHandler
 from jupyter_server.utils import url_path_join
+from .utils import parse_range
 import tornado
 import json
 import traceback
+import logging
 
-from .file_manager import FileSystemManager
-from .utils import parse_range
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 
 class BaseFileSystemHandler(APIHandler):
@@ -155,7 +158,7 @@ def initialize(self, fs_manager):
         self.fs_manager = fs_manager
 
     # POST /jupyter_fsspec/files/action?key=my-key&item_path=/some_directory/file.txt
-    def post(self):
+    async def post(self):
         """Upload/Download the resource at the input path to destination path.
 
         :param [key]: [Query arg string used to retrieve the appropriate filesystem instance]
@@ -166,32 +169,39 @@ def post(self):
         :return: dict with a status, description and (optionally) error
         :rtype: dict
         """
-        key = self.get_argument("key")
         request_data = json.loads(self.request.body.decode("utf-8"))
-        req_item_path = request_data.get("item_path")
         action = request_data.get("action")
-        destination = request_data.get("content")
-        # dest_fs_key = request_data.get('destination_key')
-        # dest_fs_info = self.fs_manager.get_filesystem(dest_fs_key)
-        # dest_path = dest_fs_info["path"]
+        local_path = request_data.get("local_path")
+        remote_path = request_data.get("remote_path")
+        dest_fs_key = request_data.get("destination_key")
+        dest_fs_info = self.fs_manager.get_filesystem(dest_fs_key)
+        dest_path = dest_fs_info["canonical_path"]
+        # if destination is subfolder, need to parse canonical_path for protocol?
 
         response = {"content": None}
 
-        fs, item_path = self.validate_fs("post", key, req_item_path)
+        fs, dest_path = self.validate_fs("post", dest_fs_key, dest_path)
         fs_instance = fs["instance"]
+        print(f"fs_instance: {fs_instance}")
 
         try:
             if action == "upload":
-                # upload     fs_instance.put(local_path, remote_path)
-                local_path = item_path
-                remote_path = destination
-                fs_instance.put(local_path, remote_path, recursive=True)
+                # upload     remote.put(local_path, remote_path)
+                logger.debug("Upload file")
+                protocol = self.fs_manager.get_filesystem_protocol(dest_fs_key)
+                if protocol not in remote_path:
+                    remote_path = protocol + remote_path
+                # TODO: handle creating directories? current: flat item upload
+                # remote_path = remote_path (root) + 'nested/'
+                await fs_instance._put(local_path, remote_path, recursive=True)
                 response["description"] = f"Uploaded {local_path} to {remote_path}."
-            else:  # download
-                # download   fs_instance.get(remote_path, local_path)
-                local_path = destination
-                remote_path = item_path
-                fs_instance.get(remote_path, local_path, recursive=True)
+            else:
+                # download   remote.get(remote_path, local_path)
+                logger.debug("Download file")
+                protocol = self.fs_manager.get_filesystem_protocol(dest_fs_key)
+                if protocol not in remote_path:
+                    remote_path = protocol + remote_path
+                await fs_instance._get(remote_path, local_path, recursive=True)
                 response["description"] = f"Downloaded {remote_path} to {local_path}."
 
             response["status"] = "success"