From 1953ecdf4dcb666b177c533571e12f69a1a3e000 Mon Sep 17 00:00:00 2001 From: Kovid Goyal Date: Mon, 23 Aug 2021 22:10:26 +0530 Subject: [PATCH] Add tests for container transmission --- kitty/file_transmission.py | 11 ++-- kitty_tests/file_transmission.py | 88 ++++++++++++++++++++++++++++++++ 2 files changed, 96 insertions(+), 3 deletions(-) diff --git a/kitty/file_transmission.py b/kitty/file_transmission.py index 4909da59f..22a627de5 100644 --- a/kitty/file_transmission.py +++ b/kitty/file_transmission.py @@ -157,10 +157,14 @@ class TarExtractor: targetpath = resolve_name(tinfo.name, dest) if targetpath is None: continue + upperdirs = os.path.dirname(targetpath) + os.makedirs(upperdirs, exist_ok=True) if tinfo.isdir(): self.tf.makedir(tinfo, targetpath) directories.append((targetpath, copy.copy(tinfo))) continue + if tinfo.islnk(): + tinfo._link_target = os.path.join(upperdirs, tinfo.linkname) # type: ignore if tinfo.isfile(): self.tf.makefile(tinfo, targetpath) elif tinfo.isfifo(): @@ -189,9 +193,8 @@ class ZipExtractor: def __call__(self, dest: str) -> None: for zinfo in self.zf.infolist(): targetpath = resolve_name(zinfo.filename, dest) - if targetpath is None: - continue - self.zf.extract(zinfo, targetpath) + if targetpath is not None: + self.zf.extract(zinfo, dest) class ActiveCommand: @@ -357,8 +360,10 @@ class FileTransmission: Container.extractor_for_container_fmt(cmd.file, cmd.ftc.container_fmt)(cmd.dest) except OSError as e: self.send_fail_on_os_error(cmd.ftc, e, 'Failed to extract files from container') + raise except Exception: self.send_response(cmd.ftc, status='EFAIL:Failed to extract files from container') + raise if not cmd.ftc.quiet: self.send_response(cmd.ftc, status='COMPLETED') finally: diff --git a/kitty_tests/file_transmission.py b/kitty_tests/file_transmission.py index c303dbcf5..b4e5d3285 100644 --- a/kitty_tests/file_transmission.py +++ b/kitty_tests/file_transmission.py @@ -5,7 +5,11 @@ import os import shutil +import tarfile import tempfile +import zipfile +import zlib +from io import BytesIO from kitty.file_transmission import ( Action, Compression, Container, FileTransmissionCommand, @@ -15,6 +19,12 @@ from kitty.file_transmission import ( from . import BaseTest +def names_in(path): + for dirpath, dirnames, filenames in os.walk(path): + for d in dirnames + filenames: + yield os.path.relpath(os.path.join(dirpath, d), path) + + def serialized_cmd(**fields) -> str: for k, A in (('action', Action), ('container_fmt', Container), ('compression', Compression)): if k in fields: @@ -35,6 +45,10 @@ class TestFileTransmission(BaseTest): def tearDown(self): shutil.rmtree(self.tdir) + def clean_tdir(self): + shutil.rmtree(self.tdir) + self.tdir = tempfile.mkdtemp() + def assertPathEqual(self, a, b): a = os.path.abspath(os.path.realpath(a)) b = os.path.abspath(os.path.realpath(b)) @@ -74,3 +88,77 @@ class TestFileTransmission(BaseTest): self.ae(ft.test_responses, [{'status': 'OK'}]) self.assertFalse(os.path.exists(dest)) self.assertFalse(ft.active_cmds) + # compress with zlib + ft = FileTransmission() + dest = os.path.join(self.tdir, '3.bin') + ft.handle_serialized_command(serialized_cmd(action='send', dest=dest, compression='zlib')) + self.ae(ft.test_responses, [{'status': 'OK'}]) + odata = 'abcd' * 1024 + data = zlib.compress(odata.encode('ascii')) + ft.handle_serialized_command(serialized_cmd(action='data', data=data[:len(data)//2])) + self.assertTrue(os.path.exists(dest)) + ft.handle_serialized_command(serialized_cmd(action='end_data', data=data[len(data)//2:])) + with open(dest) as f: + self.ae(f.read(), odata) + self.ae(ft.test_responses, [{'status': 'OK'}, {'status': 'COMPLETED'}]) + del odata + del data + + # zip send + self.clean_tdir() + buf = BytesIO() + with zipfile.ZipFile(buf, 'w') as zf: + zf.writestr('one.txt', '1' * 1111) + zf.writestr('two/one', '2' * 2222) + zf.writestr('onex/../../three', '3333') + zf.writestr('/onex', '3333') + ft = FileTransmission() + dest = os.path.join(self.tdir, 'zf') + ft.handle_serialized_command(serialized_cmd(action='send', dest=dest, container_fmt='zip')) + self.ae(ft.test_responses, [{'status': 'OK'}]) + ft.handle_serialized_command(serialized_cmd(action='end_data', data=buf.getvalue())) + self.ae(ft.test_responses, [{'status': 'OK'}, {'status': 'COMPLETED'}]) + with open(os.path.join(dest, 'one.txt')) as f: + self.ae(f.read(), '1' * 1111) + with open(os.path.join(dest, 'two', 'one')) as f: + self.ae(f.read(), '2' * 2222) + self.ae({'zf', 'zf/two', 'zf/one.txt', 'zf/two/one'}, set(names_in(self.tdir))) + + # tar send + for mode in ('', 'gz', 'bz2', 'xz'): + buf = BytesIO() + with tarfile.open(fileobj=buf, mode=f'w:{mode}') as tf: + def a(name, data, mode=0o717, lt=None): + ti = tarfile.TarInfo(name) + ti.mtime = 13 + ti.size = len(data) + ti.mode = mode + if lt: + ti.linkname = data + ti.type = lt + tf.addfile(ti) + else: + tf.addfile(ti, BytesIO(data.encode('utf-8'))) + a('a.txt', 'abcd') + a('/b.txt', 'abcd') + a('../c.txt', 'abcd') + a('sym', 'a.txt', lt=tarfile.SYMTYPE) + a('asym', '/abstarget', lt=tarfile.SYMTYPE) + a('link', 'a.txt', lt=tarfile.LNKTYPE) + self.clean_tdir() + ft = FileTransmission() + dest = os.path.join(self.tdir, 'tf') + ft.handle_serialized_command(serialized_cmd(action='send', dest=dest, container_fmt='t' + (mode or 'ar'))) + self.ae(ft.test_responses, [{'status': 'OK'}]) + ft.handle_serialized_command(serialized_cmd(action='end_data', data=buf.getvalue())) + self.ae(ft.test_responses, [{'status': 'OK'}, {'status': 'COMPLETED'}]) + with open(os.path.join(dest, 'a.txt')) as f: + self.ae(f.read(), 'abcd') + st = os.stat(f.name) + self.ae(st.st_mode & 0b111111111, 0o717) + self.ae(st.st_mtime, 13) + self.ae(os.path.realpath(os.path.join(dest, 'sym')), f.name) + self.ae(os.path.realpath(os.path.join(dest, 'asym')), '/abstarget') + self.assertTrue(os.path.samefile(f.name, os.path.join(dest, 'link'))) + self.ae({'tf', 'tf/a.txt', 'tf/sym', 'tf/asym', 'tf/link'}, set(names_in(self.tdir))) + self.ae(len(os.listdir(self.tdir)), 1)