From 8c6dc03c20e3f69cf97e3c5dd5eac01404560704 Mon Sep 17 00:00:00 2001
From: ssyssy <ssyssy@umich.edu>
Date: Mon, 29 Apr 2024 12:46:47 -0700
Subject: [PATCH] Add pyarrow usage test case (#38)

---
 setup.py                              |  1 +
 tests/fs/test_docker_pyarrow_usage.py | 47 +++++++++++++++++++++++++++
 2 files changed, 48 insertions(+)
 create mode 100644 tests/fs/test_docker_pyarrow_usage.py

diff --git a/setup.py b/setup.py
index 2bbfc1d..8864622 100644
--- a/setup.py
+++ b/setup.py
@@ -27,6 +27,7 @@
             "pytest",
             "pytest-aiohttp",
             "ray",
+            "pyarrow",
         ]
     },
     python_requires=">=3.8",
diff --git a/tests/fs/test_docker_pyarrow_usage.py b/tests/fs/test_docker_pyarrow_usage.py
new file mode 100644
index 0000000..4fd4602
--- /dev/null
+++ b/tests/fs/test_docker_pyarrow_usage.py
@@ -0,0 +1,47 @@
+import logging
+import os
+
+from pyarrow.fs import FSSpecHandler
+from pyarrow.fs import PyFileSystem
+
+from alluxiofs import AlluxioFileSystem
+from tests.conftest import ALLUXIO_FILE_PATH
+from tests.conftest import LOCAL_FILE_PATH
+from tests.fs.test_docker_fsspec_cat import ALLUXIO_PREFIX
+from tests.fs.test_docker_fsspec_cat import FILE_PATH
+
+LOGGER = logging.getLogger(__name__)
+
+
+def alluxio_pyarrow_test(py_fs, alluxio_path, local_path):
+    file_size = os.path.getsize(local_path)
+
+    file_info = py_fs.get_file_info(alluxio_path)
+    assert file_info.is_file
+    assert file_info.size == file_size
+    assert file_info.path == alluxio_path
+
+    with py_fs.open_input_file(alluxio_path) as f:
+        alluxio_file_data = f.read()
+
+    with open(local_path, "rb") as local_file:
+        local_file_data = local_file.read()
+    assert local_file_data == alluxio_file_data
+
+
+def test_alluxio_pyarrow(alluxio_file_system: AlluxioFileSystem):
+    py_fs = PyFileSystem(FSSpecHandler(alluxio_file_system))
+
+    alluxio_pyarrow_test(py_fs, ALLUXIO_FILE_PATH, LOCAL_FILE_PATH)
+    alluxio_pyarrow_test(
+        py_fs,
+        ALLUXIO_PREFIX + ALLUXIO_FILE_PATH,
+        LOCAL_FILE_PATH,
+    )
+    alluxio_pyarrow_test(py_fs, FILE_PATH, LOCAL_FILE_PATH)
+
+
+def test_etcd_alluxio_pyarrow(
+    etcd_alluxio_file_system: AlluxioFileSystem,
+):
+    test_alluxio_pyarrow(etcd_alluxio_file_system)