Add start of pyarrow IO tests

python/pyarrow/tests/test_io.py

@@ -0,0 +1,94 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from io import BytesIO
+from os.path import join as pjoin
+import os
+
+import pytest
+
+import pyarrow.io as io
+
+#----------------------------------------------------------------------
+# HDFS tests
+
+
+def hdfs_test_client():
+ host = os.environ.get('ARROW_HDFS_TEST_HOST', 'localhost')
+ user = os.environ['ARROW_HDFS_TEST_USER']
+ port = int(os.environ['ARROW_HDFS_TEST_PORT'])
+ return io.HDFSClient.connect(host, port, user)
+
+
+@pytest.fixture(scope='session')
+def hdfs():
+ return hdfs_test_client()
+
+
+HDFS_TMP_PATH = '/tmp/pyarrow-test'
+
+
+def test_hdfs_close():
+ client = hdfs_test_client()
+ assert client.is_open
+ client.close()
+ assert not client.is_open
+
+ with pytest.raises(Exception):
+ client.ls('/')
+
+
+def test_hdfs_mkdir(hdfs):
+ path = pjoin(HDFS_TMP_PATH, 'test-dir/test-dir')
+ parent_path = pjoin(HDFS_TMP_PATH, 'test-dir')
+
+ hdfs.mkdir(path)
+ assert hdfs.exists(path)
+
+ hdfs.delete(parent_path, recursive=True)
+ assert not hdfs.exists(path)
+
+
+def test_hdfs_ls(hdfs):
+ base_path = pjoin(HDFS_TMP_PATH, 'ls-test')
+ hdfs.mkdir(base_path)
+
+ dir_path = pjoin(base_path, 'a-dir')
+ f1_path = pjoin(base_path, 'a-file-1')
+
+ hdfs.mkdir(dir_path)
+
+ f = hdfs.open(f1_path, 'wb')
+ f.write('a' * 10)
+
+ contents = sorted(hdfs.ls(base_path, False))
+ assert contents == [dir_path, f1_path]
+
+
+def test_hdfs_download_upload(hdfs):
+ base_path = pjoin(HDFS_TMP_PATH, 'upload-test')
+
+ data = b'foobarbaz'
+ buf = BytesIO(data)
+ buf.seek(0)
+
+ hdfs.upload(base_path, buf)
+
+ out_buf = BytesIO()
+ hdfs.download(base_path, out_buf)
+ out_buf.seek(0)
+ assert out_buf.getvalue() == data