Traktor/myenv/Lib/site-packages/fsspec/implementations/tests/test_sftp.py
2024-05-23 01:57:24 +02:00

234 lines
6.7 KiB
Python

import os
import shlex
import subprocess
import time
from tarfile import TarFile
import pytest
import fsspec
pytest.importorskip("paramiko")
def stop_docker(name):
cmd = shlex.split(f'docker ps -a -q --filter "name={name}"')
cid = subprocess.check_output(cmd).strip().decode()
if cid:
subprocess.call(["docker", "rm", "-f", cid])
@pytest.fixture(scope="module")
def ssh():
try:
pchk = ["docker", "run", "--name", "fsspec_test_sftp", "hello-world"]
subprocess.check_call(pchk)
stop_docker("fsspec_test_sftp")
except (subprocess.CalledProcessError, FileNotFoundError):
pytest.skip("docker run not available")
return
# requires docker
cmds = [
r"apt-get update",
r"apt-get install -y openssh-server",
r"mkdir /var/run/sshd",
"bash -c \"echo 'root:pass' | chpasswd\"",
(
r"sed -i 's/PermitRootLogin prohibit-password/PermitRootLogin yes/' "
r"/etc/ssh/sshd_config"
),
(
r"sed 's@session\s*required\s*pam_loginuid.so@session optional "
r"pam_loginuid.so@g' -i /etc/pam.d/sshd"
),
r'bash -c "echo \"export VISIBLE=now\" >> /etc/profile"',
r"/usr/sbin/sshd",
]
name = "fsspec_sftp"
stop_docker(name)
cmd = f"docker run -d -p 9200:22 --name {name} ubuntu:16.04 sleep 9000"
try:
cid = subprocess.check_output(shlex.split(cmd)).strip().decode()
for cmd in cmds:
subprocess.call(["docker", "exec", cid] + shlex.split(cmd))
time.sleep(1)
yield {
"host": "localhost",
"port": 9200,
"username": "root",
"password": "pass",
}
finally:
stop_docker(name)
@pytest.fixture(scope="module")
def root_path():
return "/home/someuser/"
def test_simple(ssh, root_path):
f = fsspec.get_filesystem_class("sftp")(**ssh)
f.mkdirs(root_path + "deeper")
try:
f.touch(root_path + "deeper/afile")
assert f.find(root_path) == [root_path + "deeper/afile"]
assert f.ls(root_path + "deeper/") == [root_path + "deeper/afile"]
assert f.info(root_path + "deeper/afile")["type"] == "file"
assert f.info(root_path + "deeper/afile")["size"] == 0
assert f.exists(root_path)
finally:
f.rm(root_path, recursive=True)
assert not f.exists(root_path)
@pytest.mark.parametrize("protocol", ["sftp", "ssh"])
def test_with_url(protocol, ssh):
fo = fsspec.open(
protocol
+ "://{username}:{password}@{host}:{port}/home/someuserout".format(**ssh),
"wb",
)
with fo as f:
f.write(b"hello")
fo = fsspec.open(
protocol
+ "://{username}:{password}@{host}:{port}/home/someuserout".format(**ssh),
"rb",
)
with fo as f:
assert f.read() == b"hello"
@pytest.mark.parametrize("protocol", ["sftp", "ssh"])
def test_get_dir(protocol, ssh, root_path, tmpdir):
path = str(tmpdir)
f = fsspec.filesystem(protocol, **ssh)
f.mkdirs(root_path + "deeper", exist_ok=True)
f.touch(root_path + "deeper/afile")
f.get(root_path, path, recursive=True)
assert os.path.isdir(f"{path}/deeper")
assert os.path.isfile(f"{path}/deeper/afile")
f.get(
protocol
+ "://{username}:{password}@{host}:{port}{root_path}".format(
root_path=root_path, **ssh
),
f"{path}/test2",
recursive=True,
)
assert os.path.isdir(f"{path}/test2/deeper")
assert os.path.isfile(f"{path}/test2/deeper/afile")
@pytest.fixture(scope="module")
def netloc(ssh):
username = ssh.get("username")
password = ssh.get("password")
host = ssh.get("host")
port = ssh.get("port")
userpass = (
f"{username}:{password if password is not None else ''}@"
if username is not None
else ""
)
netloc = f"{host}:{port if port is not None else ''}"
return userpass + netloc
def test_put_file(ssh, tmp_path, root_path):
tmp_file = tmp_path / "a.txt"
with open(tmp_file, mode="w") as fd:
fd.write("blabla")
f = fsspec.get_filesystem_class("sftp")(**ssh)
f.put_file(lpath=tmp_file, rpath=root_path + "a.txt")
def test_simple_with_tar(ssh, netloc, tmp_path, root_path):
files_to_pack = ["a.txt", "b.txt"]
tar_filename = make_tarfile(files_to_pack, tmp_path)
f = fsspec.get_filesystem_class("sftp")(**ssh)
f.mkdirs(f"{root_path}deeper", exist_ok=True)
try:
remote_tar_filename = f"{root_path}deeper/somefile.tar"
with f.open(remote_tar_filename, mode="wb") as wfd:
with open(tar_filename, mode="rb") as rfd:
wfd.write(rfd.read())
fs = fsspec.open(f"tar::ssh://{netloc}{remote_tar_filename}").fs
files = fs.find("/")
assert files == files_to_pack
finally:
f.rm(root_path, recursive=True)
def make_tarfile(files_to_pack, tmp_path):
"""Create a tarfile with some files."""
tar_filename = tmp_path / "sometarfile.tar"
for filename in files_to_pack:
with open(tmp_path / filename, mode="w") as fd:
fd.write("")
with TarFile(tar_filename, mode="w") as tf:
for filename in files_to_pack:
tf.add(tmp_path / filename, arcname=filename)
return tar_filename
def test_transaction(ssh, root_path):
f = fsspec.get_filesystem_class("sftp")(**ssh)
f.mkdirs(root_path + "deeper", exist_ok=True)
try:
f.start_transaction()
f.touch(root_path + "deeper/afile")
assert f.find(root_path) == []
f.end_transaction()
assert f.find(root_path) == [root_path + "deeper/afile"]
with f.transaction:
assert f._intrans
f.touch(root_path + "deeper/afile2")
assert f.find(root_path) == [root_path + "deeper/afile"]
assert f.find(root_path) == [
root_path + "deeper/afile",
root_path + "deeper/afile2",
]
finally:
f.rm(root_path, recursive=True)
@pytest.mark.parametrize("path", ["/a/b/c", "a/b/c"])
def test_mkdir_create_parent(ssh, path):
f = fsspec.get_filesystem_class("sftp")(**ssh)
with pytest.raises(FileNotFoundError):
f.mkdir(path, create_parents=False)
f.mkdir(path)
assert f.exists(path)
with pytest.raises(FileExistsError, match=path):
f.mkdir(path)
f.rm(path, recursive=True)
assert not f.exists(path)
@pytest.mark.parametrize("path", ["/a/b/c", "a/b/c"])
def test_makedirs_exist_ok(ssh, path):
f = fsspec.get_filesystem_class("sftp")(**ssh)
f.makedirs(path, exist_ok=False)
with pytest.raises(FileExistsError, match=path):
f.makedirs(path, exist_ok=False)
f.makedirs(path, exist_ok=True)
f.rm(path, recursive=True)
assert not f.exists(path)