def test_walk(self, local_mock_dir, mock_data_dir):
"""Tests the `walk` method."""
with LocalHook() as hook:
entries = list(hook.walk(local_mock_dir))
pytest.helpers.assert_walk_equal(entries, os.walk(mock_data_dir))
def test_postgres_to_local_csv(mocker, postgres, tmp_path):
mocker.patch.object(
PostgresHook,
"get_connection",
return_value=Connection(
conn_id="test",
conn_type="postgres",
host="localhost",
login="******",
password="******",
port=postgres.ports["5432/tcp"][0],
),
)
task = DbToFsOperator(
task_id="test_id",
src_db_hook=PostgresHook(conn_id="this_conn_will_be_mocked"),
src_query="SELECT * FROM dummy",
output_filetype="csv",
dest_path=str(tmp_path / "test.csv"),
dest_fs_hook=LocalHook(),
)
task.execute(context={})
# Assert output
with open(os.path.join(os.path.dirname(__file__),
"testdata.csv")) as local_file:
csv_reader = csv.reader(local_file)
local_data = list(csv_reader)
with (tmp_path / "test.csv").open(mode="r") as result_file:
csv_reader = csv.reader(result_file)
result_data = list(csv_reader)
assert local_data == result_data
def test_glob(self, local_mock_dir, mock_data_dir):
"""Tests the `glob` method."""
with LocalHook() as hook:
# Test simple glob on txt files.
txt_files = hook.glob(posixpath.join(local_mock_dir, "*.txt"))
txt_expected = glob.glob(os.path.join(mock_data_dir, "*.txt"))
assert_paths_equal(txt_files,
txt_expected,
root_a=local_mock_dir,
root_b=mock_data_dir)
# Test glob for non-existing tsv files.
assert not set(hook.glob(posixpath.join(local_mock_dir, "*.tsv")))
# Test glob on directory.
dir_paths = hook.glob(posixpath.join(local_mock_dir, "subdir"))
assert set(strip_root(p, root=local_mock_dir)
for p in dir_paths) == {"subdir"}
# Test glob with dir pattern.
file_paths = hook.glob(posixpath.join(local_mock_dir, "*",
"*.txt"))
expected_paths = glob.glob(
os.path.join(mock_data_dir, "*", "*.txt"))
assert_paths_equal(file_paths,
expected_paths,
root_a=local_mock_dir,
root_b=mock_data_dir)
def test_files_missing(self, local_mock_dir, test_dag):
"""Tests example with missing files."""
task = sensors.FileSensor(path=posixpath.join(local_mock_dir, "*.xml"),
hook=LocalHook(),
task_id="file_sensor",
dag=test_dag)
assert not task.poke({})
def test_files_present(self, local_mock_dir, test_dag):
"""Tests example with present files."""
task = sensors.FileSensor(path=posixpath.join(local_mock_dir, "*.txt"),
hook=LocalHook(),
task_id="file_sensor",
dag=test_dag)
assert task.poke({})
def test_exists(self, local_mock_dir):
"""Tests the `exists` method."""
with LocalHook() as hook:
assert hook.exists(posixpath.join(local_mock_dir, "subdir"))
assert hook.exists(posixpath.join(local_mock_dir, "test.txt"))
assert not hook.exists(
posixpath.join(local_mock_dir, "non-existing.txt"))
def test_rmtree(self, local_mock_dir):
"""Tests the `rmtree` method."""
dir_path = posixpath.join(local_mock_dir, "subdir")
assert posixpath.exists(dir_path)
with LocalHook() as hook:
hook.rmtree(dir_path)
assert not posixpath.exists(dir_path)
def test_rm(self, local_mock_dir):
"""Tests the `rm` method."""
file_path = posixpath.join(local_mock_dir, "test.txt")
assert posixpath.exists(file_path)
with LocalHook() as hook:
hook.rm(file_path)
assert not posixpath.exists(file_path)
def test_open_read(self, local_mock_dir):
"""Tests reading of a file using the `open` method."""
file_path = posixpath.join(local_mock_dir, "test.txt")
with LocalHook() as hook:
with hook.open(file_path) as file_:
content = file_.read()
assert content == b"Test file\n"
def test_makedirs(self, tmpdir):
"""Tests the `mkdir` method with mode parameter."""
dir_path = posixpath.join(str(tmpdir), "some", "nested", "dir")
with LocalHook() as hook:
hook.makedirs(dir_path, mode=0o750)
assert posixpath.exists(dir_path)
assert oct(os.stat(dir_path).st_mode)[-3:] == "750"
def test_non_existing_dir(self, local_mock_dir, test_dag):
"""Tests example with non-existing dir."""
task = sensors.FileSensor(path=posixpath.join(local_mock_dir,
"non_existing_dir",
"*.xml"),
hook=LocalHook(),
task_id="file_sensor",
dag=test_dag)
assert not task.poke({})
def test_open_write(self, tmpdir):
"""Tests writing of a file using the `open` method."""
file_path = posixpath.join(str(tmpdir), "test2.txt")
assert not posixpath.exists(file_path)
with LocalHook() as hook:
with hook.open(file_path, "wb") as file_:
file_.write(b"Test file\n")
assert posixpath.exists(file_path)
def test_makedirs_exists(self, tmpdir):
"""Tests the `mkdir` method with exists_ok parameter."""
dir_path = posixpath.join(str(tmpdir), "some", "nested", "dir")
with LocalHook() as hook:
hook.makedirs(dir_path, exist_ok=False)
with pytest.raises(IOError):
hook.makedirs(dir_path, exist_ok=False)
hook.makedirs(dir_path, exist_ok=True)
def test_mkdir_exists(self, tmpdir):
"""Tests the `mkdir` method with the exists_ok parameter."""
dir_path = posixpath.join(str(tmpdir), "subdir")
assert not posixpath.exists(dir_path)
with LocalHook() as hook:
hook.mkdir(dir_path, exist_ok=False)
with pytest.raises(IOError):
hook.mkdir(dir_path, exist_ok=False)
hook.mkdir(dir_path, exist_ok=True)
def test_glob_recursive(self, local_mock_dir, mock_data_dir):
"""Tests the `glob` method with recursive = True."""
with LocalHook() as hook:
file_paths = hook.glob(posixpath.join(local_mock_dir, "**",
"*.txt"),
recursive=True)
expected_paths = set(
glob.glob(os.path.join(mock_data_dir, "**", "*.txt"),
recursive=True))
assert_paths_equal(file_paths,
expected_paths,
root_a=local_mock_dir,
root_b=mock_data_dir)
def __init__(self,
src_db_hook: DbHook,
src_query: str,
output_filetype: str,
dest_path: str,
dest_fs_hook: FsHook = None,
**kwargs):
super().__init__(**kwargs)
self._src_db_hook = src_db_hook
self._src_query = src_query
self._output_filetype = output_filetype
self._dest_path = dest_path
self._dest_fs_hook = dest_fs_hook or LocalHook()
def test_postgres_to_local_json(mocker, postgres, tmp_path):
mocker.patch.object(
PostgresHook,
"get_connection",
return_value=Connection(
conn_id="test",
conn_type="postgres",
host="localhost",
login="******",
password="******",
port=postgres.ports["5432/tcp"][0],
),
)
task = DbToFsOperator(
task_id="test_id",
src_db_hook=PostgresHook(conn_id="this_conn_will_be_mocked"),
src_query="SELECT * FROM dummy",
output_filetype="json",
dest_path=str(tmp_path / "test.json"),
dest_fs_hook=LocalHook(),
)
task.execute(context={})
# Assert output
with open(os.path.join(os.path.dirname(__file__),
"testdata.csv")) as local_file:
csv_reader = csv.reader(local_file)
header = next(csv_reader)
rows = list(csv_reader)
local_data = [{header[i]: col
for i, col in enumerate(row)}
for row in rows] # convert to dict
with (tmp_path / "test.json").open(mode="r") as result_file:
result_data = json.load(result_file)
# CSV reader makes everything string, so cast all columns to string for comparing values
result_data = [{str(k): str(v)
for k, v in row.items()} for row in result_data]
assert local_data == result_data
def __init__(self, path, hook=None, **kwargs):
super(FileSensor, self).__init__(**kwargs)
self._path = path
self._hook = hook or LocalHook()
def test_isdir(self, local_mock_dir):
"""Tests the `isdir` method."""
with LocalHook() as hook:
assert hook.isdir(posixpath.join(local_mock_dir, "subdir"))
assert not hook.isdir(posixpath.join(local_mock_dir, "test.txt"))
def test_listdir(self, local_mock_dir, mock_data_dir):
"""Tests the `listdir` method."""
with LocalHook() as hook:
assert set(hook.listdir(local_mock_dir)) == set(
os.listdir(mock_data_dir))