Clean up, remove compat with py < 3.6

[pyi2ncommon] / test / test_zip_stream.py

# The software in this package is distributed under the GNU General
# Public License version 2 (with a special exception described below).
#
# A copy of GNU General Public License (GPL) is included in this distribution,
# in the file COPYING.GPL.
#
# As a special exception, if other files instantiate templates or use macros
# or inline functions from this file, or you compile this file and link it
# with other works to produce a work based on this file, this file
# does not by itself cause the resulting work to be covered
# by the GNU General Public License.
#
# However the source code for this file must still be made available
# in accordance with section (3) of the GNU General Public License.
#
# This exception does not invalidate any other reasons why a work based
# on this file might be covered by the GNU General Public License.

""" test_zip_stream.py: unit tests for zip_stream

Test classes and functions in :py:mod:`pyi2ncommon.zip_stream`.

.. codeauthor:: Intra2net
"""

import unittest
from tempfile import mkdtemp
import os
from os.path import join, isfile
import shutil
import io
import zipfile       # just to make sure: for comparing read zips using zipfile

# relative import of tested module ensures we do not test installed version
try:
    from src.zip_stream import ZipStream
except ImportError as ie:
    # prevent misleading error from unittest loader
    # AttributeError: 'module' object has no attribute 'test_zip_stream'
    raise RuntimeError('Failed to import tested module: {}'.format(ie))


# test data
TEXT_DATA = 'Test text file\nfor testing pyi2ncommon.zip_stream\n'
BIN_DATA = b'Test binary file\nfor testing pyi2ncommon.zip_stream\n' \
           b'Bytes: \x00\x01\x02\nEnd\n'

# test file names (all in temp test dir)
TEXT_FILE = 'text.txt'
BIN_FILE = 'binary.bin'
SUBDIR = 'subdir'
ZIP_FILE = 'test.zip'


class BytesIONoSeekNorRead(io.BytesIO):
    """Subclass of :py:class:`io.BytesIO` with seek() and read() removed"""
    def seekable(self):
        return False

    def readable(self):
        return False

    def seek(self, *args):
        raise AttributeError('this function was removed')

    def readline(self, *args):
        raise AttributeError('this function was removed')

    def readlines(self, *args):
        raise AttributeError('this function was removed')

    def read(self, *args):
        raise AttributeError('this function was removed')

    def readinto(self, *args):
        raise AttributeError('this function was removed')

    def readall(self, *args):
        raise AttributeError('this function was removed')

    def read1(self, *args):
        raise AttributeError('this function was removed')


class ZipStreamTester(unittest.TestCase):
    """ only test case in this module, see module doc for more help """

    #: directory used for all files; managed in setUpClass / tearDownClass
    temp_dir = None

    @classmethod
    def temp_path(cls, *path_components):
        """Quick alias to create file name in temp dir"""
        return join(cls.temp_dir, *path_components)

    @classmethod
    def setUpClass(cls):
        """
        called once before tests in this class

        creates temp dir with a few files for zipping in it
        """
        cls.temp_dir = mkdtemp(prefix='pyi2ncommon-test-zip-stream-')
        with open(cls.temp_path(TEXT_FILE), 'wt') as writer:
            writer.write(TEXT_DATA)
        with open(cls.temp_path(BIN_FILE), 'wb') as writer:
            writer.write(BIN_DATA)
        os.mkdir(cls.temp_path(SUBDIR))

    @classmethod
    def tearDownClass(cls):
        """called once when all tests in this class are done. rm temp_dir"""
        shutil.rmtree(cls.temp_dir)

    def tearDown(self):
        """ called after each test function """
        if isfile(self.temp_path(ZIP_FILE)):
            os.unlink(self.temp_path(ZIP_FILE))

    def _test_zip(self, test_subdir=False):
        """Helper for test_* functions: check given zip for contents"""
        expect_contents = {TEXT_FILE, BIN_FILE}
        if test_subdir:
            expect_contents.add(SUBDIR + '/')

        with zipfile.ZipFile(self.temp_path(ZIP_FILE), 'r') as unzipper:
            self.assertEqual(unzipper.testzip(), None)
            self.assertSetEqual(set(unzipper.namelist()), expect_contents)
            with unzipper.open(TEXT_FILE) as reader:
                self.assertEqual(reader.read().decode('utf8'), TEXT_DATA)
            with unzipper.open(BIN_FILE) as reader:
                self.assertEqual(reader.read(), BIN_DATA)

    def test_base_class(self):
        """Check that we did not destroy ZipStream's base class's behaviour"""
        # create zip archive using functions in ZipFile
        with ZipStream(self.temp_path(ZIP_FILE), 'w') as zipper:
            zipper.write(self.temp_path(TEXT_FILE), TEXT_FILE)
            zipper.writestr(BIN_FILE, BIN_DATA)
            zipper.write(self.temp_path(SUBDIR), SUBDIR)
        self._test_zip(test_subdir=True)

    def test_stream_input(self):
        """Test function write_stream we added; write to file"""
        # create zip archive using write_stream
        with ZipStream(self.temp_path(ZIP_FILE), 'w') as zipper:
            info = zipper.create_zipinfo(self.temp_path(BIN_FILE), BIN_FILE)
            with open(self.temp_path(BIN_FILE), 'rb') as reader:
                zipper.write_stream(reader, info)
            info = zipper.create_zipinfo(self.temp_path(TEXT_FILE), TEXT_FILE)
            with open(self.temp_path(TEXT_FILE), 'rb') as reader:  # read byte!
                zipper.write_stream(reader, info)
        self._test_zip()

    def test_stream_output(self):
        """Test writing to an output stream like sys.stdout"""
        # create zip archive that writes to buffer
        output = BytesIONoSeekNorRead()
        with ZipStream(output, 'w') as zipper:
            info = zipper.create_zipinfo(self.temp_path(BIN_FILE), BIN_FILE)
            with open(self.temp_path(BIN_FILE), 'rb') as reader:
                zipper.write_stream(reader, info)
            info = zipper.create_zipinfo(self.temp_path(TEXT_FILE), TEXT_FILE)
            with open(self.temp_path(TEXT_FILE), 'rb') as reader:  # read byte!
                zipper.write_stream(reader, info)

        # now write to file and test
        with open(self.temp_path(ZIP_FILE), 'wb') as writer:
            writer.write(output.getvalue())
        self._test_zip()


if __name__ == '__main__':
    unittest.main()