# -*- coding: utf-8 -*-

#

# 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.

#

"""Unit tests for LocalFileSystem."""

import contextlib

import filecmp

import logging

import os

import shutil

import tempfile

import unittest

import mock

from parameterized import param

from parameterized import parameterized

from apache_beam.io import localfilesystem

from apache_beam.io.filesystem import BeamIOError

from apache_beam.options.pipeline_options import PipelineOptions

# pytype: skip-file

def _gen_fake_join(separator):

"""Returns a callable that joins paths with the given separator."""

def _join(first_path, *paths):

return separator.join((first_path.rstrip(separator), ) + paths)

return _join

def _gen_fake_split(separator):

"""Returns a callable that splits a with the given separator."""

def _split(path):

sep_index = path.rfind(separator)

if sep_index >= 0:

return (path[:sep_index], path[sep_index + 1:])

else:

return (path, '')

return _split

class LocalFileSystemTest(unittest.TestCase):

def setUp(self):

self.tmpdir = tempfile.mkdtemp()

pipeline_options = PipelineOptions()

self.fs = localfilesystem.LocalFileSystem(pipeline_options)

def tearDown(self):

shutil.rmtree(self.tmpdir)

@contextlib.contextmanager

def tmpdir_as_cwd(self):

"""Context manager that sets the current working directory to a temp dir."""

old_cwd = os.getcwd()

os.chdir(self.tmpdir)

try:

yield

finally:

os.chdir(old_cwd)

def test_scheme(self):

self.assertIsNone(self.fs.scheme())

self.assertIsNone(localfilesystem.LocalFileSystem.scheme())

def test_create_cwd_file(self):

with self.tmpdir_as_cwd():

with self.fs.create("blah.txt") as f:

f.write(b"blah")

with open("blah.txt", "rb") as f:

assert f.read() == b"blah"

@mock.patch('apache_beam.io.localfilesystem.os')

def test_unix_path_join(self, *unused_mocks):

# Test joining of Unix paths.

localfilesystem.os.path.join.side_effect = _gen_fake_join('/')

self.assertEqual(

'/tmp/path/to/file', self.fs.join('/tmp/path', 'to', 'file'))

self.assertEqual('/tmp/path/to/file', self.fs.join('/tmp/path', 'to/file'))

@mock.patch('apache_beam.io.localfilesystem.os')

def test_windows_path_join(self, *unused_mocks):

# Test joining of Windows paths.

localfilesystem.os.path.join.side_effect = _gen_fake_join('\\')

self.assertEqual(

r'C:\tmp\path\to\file', self.fs.join(r'C:\tmp\path', 'to', 'file'))

self.assertEqual(

r'C:\tmp\path\to\file', self.fs.join(r'C:\tmp\path', r'to\file'))

@mock.patch('apache_beam.io.localfilesystem.os')

def test_unix_path_split(self, os_mock):

os_mock.path.abspath.side_effect = lambda a: a

os_mock.path.split.side_effect = _gen_fake_split('/')

self.assertEqual(('/tmp/path/to', 'file'),

self.fs.split('/tmp/path/to/file'))

# Actual os.path.split will split following to '/' and 'tmp' when run in

# Unix.

self.assertEqual(('', 'tmp'), self.fs.split('/tmp'))

@mock.patch('apache_beam.io.localfilesystem.os')

def test_windows_path_split(self, os_mock):

os_mock.path.abspath = lambda a: a

os_mock.path.split.side_effect = _gen_fake_split('\\')

self.assertEqual((r'C:\tmp\path\to', 'file'),

self.fs.split(r'C:\tmp\path\to\file'))

# Actual os.path.split will split following to 'C:\' and 'tmp' when run in

# Windows.

self.assertEqual((r'C:', 'tmp'), self.fs.split(r'C:\tmp'))

def test_mkdirs(self):

path = os.path.join(self.tmpdir, 't1/t2')

self.fs.mkdirs(path)

self.assertTrue(os.path.isdir(path))

def test_mkdirs_failed(self):

path = os.path.join(self.tmpdir, 't1/t2')

self.fs.mkdirs(path)

# Check IOError if existing directory is created

with self.assertRaises(IOError):

self.fs.mkdirs(path)

with self.assertRaises(IOError):

self.fs.mkdirs(os.path.join(self.tmpdir, 't1'))

def test_match_file(self):

path = os.path.join(self.tmpdir, 'f1')

open(path, 'a').close()

# Match files in the temp directory

result = self.fs.match([path])[0]

files = [f.path for f in result.metadata_list]

self.assertEqual(files, [path])

def test_match_file_empty(self):

path = os.path.join(self.tmpdir, 'f2') # Does not exist

# Match files in the temp directory

result = self.fs.match([path])[0]

files = [f.path for f in result.metadata_list]

self.assertEqual(files, [])

def test_match_file_exception(self):

# Match files with None so that it throws an exception

with self.assertRaisesRegex(BeamIOError,

r'^Match operation failed') as error:

self.fs.match([None])

self.assertEqual(list(error.exception.exception_details.keys()), [None])

@parameterized.expand([

param('*', files=['a', 'b', os.path.join('c', 'x')], expected=['a', 'b']),

param(

'**',

files=['a', os.path.join('b', 'x'), os.path.join('c', 'x')],

expected=['a', os.path.join('b', 'x'), os.path.join('c', 'x')]),

param(

os.path.join('*', '*'),

files=[

'a',

os.path.join('b', 'x'),

os.path.join('c', 'x'),

os.path.join('d', 'x', 'y')

],

expected=[os.path.join('b', 'x'), os.path.join('c', 'x')]),

param(

os.path.join('**', '*'),

files=[

'a',

os.path.join('b', 'x'),

os.path.join('c', 'x'),

os.path.join('d', 'x', 'y')

],

expected=[

os.path.join('b', 'x'),

os.path.join('c', 'x'),

os.path.join('d', 'x', 'y')

]),

])

def test_match_glob(self, pattern, files, expected):

for filename in files:

full_path = os.path.join(self.tmpdir, filename)

dirname = os.path.dirname(full_path)

if not dirname == full_path:

# Make sure we don't go outside the tmpdir

assert os.path.commonprefix([self.tmpdir, full_path]) == self.tmpdir

try:

self.fs.mkdirs(dirname)

except IOError:

# Directory exists

pass

open(full_path, 'a').close() # create empty file

# Match both the files in the directory

full_pattern = os.path.join(self.tmpdir, pattern)

result = self.fs.match([full_pattern])[0]

files = [os.path.relpath(f.path, self.tmpdir) for f in result.metadata_list]

self.assertCountEqual(files, expected)

def test_match_directory(self):

result = self.fs.match([self.tmpdir])[0]

files = [f.path for f in result.metadata_list]

self.assertEqual(files, [self.tmpdir])

def test_match_directory_contents(self):

path1 = os.path.join(self.tmpdir, 'f1')

path2 = os.path.join(self.tmpdir, 'f2')

open(path1, 'a').close()

open(path2, 'a').close()

result = self.fs.match([os.path.join(self.tmpdir, '*')])[0]

files = [f.path for f in result.metadata_list]

self.assertCountEqual(files, [path1, path2])

def test_copy(self):

path1 = os.path.join(self.tmpdir, 'f1')

path2 = os.path.join(self.tmpdir, 'f2')

with open(path1, 'a') as f:

f.write('Hello')

self.fs.copy([path1], [path2])

self.assertTrue(filecmp.cmp(path1, path2))

def test_copy_error(self):

path1 = os.path.join(self.tmpdir, 'f1')

path2 = os.path.join(self.tmpdir, 'f2')

with self.assertRaisesRegex(BeamIOError,

r'^Copy operation failed') as error:

self.fs.copy([path1], [path2])

self.assertEqual(

list(error.exception.exception_details.keys()), [(path1, path2)])

def test_copy_directory(self):

path_t1 = os.path.join(self.tmpdir, 't1')

path_t2 = os.path.join(self.tmpdir, 't2')

self.fs.mkdirs(path_t1)

self.fs.mkdirs(path_t2)

path1 = os.path.join(path_t1, 'f1')

path2 = os.path.join(path_t2, 'f1')

with open(path1, 'a') as f:

f.write('Hello')

self.fs.copy([path_t1], [path_t2])

self.assertTrue(filecmp.cmp(path1, path2))

def test_rename(self):

path1 = os.path.join(self.tmpdir, 'f1')

path2 = os.path.join(self.tmpdir, 'f2')

with open(path1, 'a') as f:

f.write('Hello')

self.fs.rename([path1], [path2])

self.assertTrue(self.fs.exists(path2))

self.assertFalse(self.fs.exists(path1))

def test_rename_error(self):

path1 = os.path.join(self.tmpdir, 'f1')

path2 = os.path.join(self.tmpdir, 'f2')

with self.assertRaisesRegex(BeamIOError,

r'^Rename operation failed') as error:

self.fs.rename([path1], [path2])

self.assertEqual(

list(error.exception.exception_details.keys()), [(path1, path2)])

def test_rename_directory(self):

path_t1 = os.path.join(self.tmpdir, 't1')

path_t2 = os.path.join(self.tmpdir, 't2')

self.fs.mkdirs(path_t1)

path1 = os.path.join(path_t1, 'f1')

path2 = os.path.join(path_t2, 'f1')

with open(path1, 'a') as f:

f.write('Hello')

self.fs.rename([path_t1], [path_t2])

self.assertTrue(self.fs.exists(path_t2))

self.assertFalse(self.fs.exists(path_t1))

self.assertTrue(self.fs.exists(path2))

self.assertFalse(self.fs.exists(path1))

def test_exists(self):

path1 = os.path.join(self.tmpdir, 'f1')

path2 = os.path.join(self.tmpdir, 'f2')

with open(path1, 'a') as f:

f.write('Hello')

self.assertTrue(self.fs.exists(path1))

self.assertFalse(self.fs.exists(path2))

def test_checksum(self):

path1 = os.path.join(self.tmpdir, 'f1')

path2 = os.path.join(self.tmpdir, 'f2')

with open(path1, 'a') as f:

f.write('Hello')

with open(path2, 'a') as f:

f.write('foo')

# tests that localfilesystem checksum returns file size

checksum1 = self.fs.checksum(path1)

checksum2 = self.fs.checksum(path2)

self.assertEqual(checksum1, str(5))

self.assertEqual(checksum2, str(3))

# tests that fs.checksum and str(fs.size) are consistent

self.assertEqual(checksum1, str(self.fs.size(path1)))

self.assertEqual(checksum2, str(self.fs.size(path2)))

def make_tree(self, path, value, expected_leaf_count=None):

"""Create a file+directory structure from a simple dict-based DSL

:param path: root path to create directories+files under

:param value: a specification of what ``path`` should contain: ``None`` to

make it an empty directory, a string literal to make it a file with those

contents, and a ``dict`` to make it a non-empty directory and recurse

:param expected_leaf_count: only be set at the top of a recursive call

stack; after the whole tree has been created, verify the presence and

number of all files+directories, as a sanity check

"""

if value is None:

# empty directory

os.makedirs(path)

elif isinstance(value, str):

# file with string-literal contents

dir = os.path.dirname(path)

if not os.path.exists(dir):

os.makedirs(dir)

with open(path, 'a') as f:

f.write(value)

elif isinstance(value, dict):

# recurse to create a subdirectory tree

for basename, v in value.items():

self.make_tree(os.path.join(path, basename), v)

else:

raise Exception('Unexpected value in tempdir tree: %s' % value)

if expected_leaf_count is not None:

self.assertEqual(self.check_tree(path, value), expected_leaf_count)

def check_tree(self, path, value, expected_leaf_count=None):

"""Verify a directory+file structure according to the rules described in

``make_tree``

:param path: path to check under

:param value: DSL-representation of expected files+directories under

``path``

:return: number of leaf files/directories that were verified

"""

actual_leaf_count = None

if value is None:

# empty directory

self.assertTrue(os.path.exists(path), msg=path)

self.assertEqual(os.listdir(path), [])

actual_leaf_count = 1

elif isinstance(value, str):

# file with string-literal contents

with open(path, 'r') as f:

self.assertEqual(f.read(), value, msg=path)

actual_leaf_count = 1

elif isinstance(value, dict):

# recurse to check subdirectory tree

actual_leaf_count = sum([

self.check_tree(os.path.join(path, basename), v)

for basename, v in value.items()

])

else:

raise Exception('Unexpected value in tempdir tree: %s' % value)

if expected_leaf_count is not None:

self.assertEqual(actual_leaf_count, expected_leaf_count)

return actual_leaf_count

_test_tree = {

'path1': '111',

'path2': {

'2': '222', 'emptydir': None

},

'aaa': {

'b1': 'b1', 'b2': None, 'bbb': {

'ccc': {

'ddd': 'DDD'

}

}, 'c': None

}

}

def test_delete_globs(self):

dir = os.path.join(self.tmpdir, 'dir')

self.make_tree(dir, self._test_tree, expected_leaf_count=7)

self.fs.delete([os.path.join(dir, 'path*'), os.path.join(dir, 'aaa', 'b*')])

# One empty nested directory is left

self.check_tree(dir, {'aaa': {'c': None}}, expected_leaf_count=1)

def test_recursive_delete(self):

dir = os.path.join(self.tmpdir, 'dir')

self.make_tree(dir, self._test_tree, expected_leaf_count=7)

self.fs.delete([dir])

self.check_tree(self.tmpdir, {'': None}, expected_leaf_count=1)

def test_delete_glob_errors(self):

dir = os.path.join(self.tmpdir, 'dir')

self.make_tree(dir, self._test_tree, expected_leaf_count=7)

with self.assertRaisesRegex(BeamIOError,

r'^Delete operation failed') as error:

self.fs.delete([

os.path.join(dir, 'path*'),

os.path.join(dir, 'aaa', 'b*'),

os.path.join(dir, 'aaa', 'd*') # doesn't match anything, will raise

])

self.check_tree(dir, {'aaa': {'c': None}}, expected_leaf_count=1)

self.assertEqual(

list(error.exception.exception_details.keys()),

[os.path.join(dir, 'aaa', 'd*')])

with self.assertRaisesRegex(BeamIOError,

r'^Delete operation failed') as error:

self.fs.delete([

os.path.join(dir, 'path*') # doesn't match anything, will raise

])

self.check_tree(dir, {'aaa': {'c': None}}, expected_leaf_count=1)

self.assertEqual(

list(error.exception.exception_details.keys()),

[os.path.join(dir, 'path*')])

def test_delete(self):

path1 = os.path.join(self.tmpdir, 'f1')

with open(path1, 'a') as f:

f.write('Hello')

self.assertTrue(self.fs.exists(path1))

self.fs.delete([path1])

self.assertFalse(self.fs.exists(path1))

def test_delete_error(self):

path1 = os.path.join(self.tmpdir, 'f1')

with self.assertRaisesRegex(BeamIOError,

r'^Delete operation failed') as error:

self.fs.delete([path1])

self.assertEqual(list(error.exception.exception_details.keys()), [path1])

if __name__ == '__main__':

logging.getLogger().setLevel(logging.INFO)

unittest.main()