~abentley/bzrtools/bzrtools.dev

import os

from StringIO import StringIO

from shutil import rmtree, copy2, copytree

import tarfile

import tempfile

from unittest import makeSuite

from bzrlib import osutils

from bzrlib.bzrdir import BzrDir

from bzrlib.plugins.bzrtools.upstream_import import (

    common_directory,

    import_archive,

    import_tar,

    import_zip,

    import_dir,

    top_directory,

    ZipFileWrapper, 

)

from bzrlib.tests import TestCaseInTempDir

def import_tar_broken(tree, tar_input):

    """

    Import a tarfile with names that that end in //, e.g. Feisty Python 2.5

    """

    tar_file = tarfile.open('lala', 'r', tar_input)

    for member in tar_file.members:

        if member.name.endswith('/'):

            member.name += '/'

    import_archive(tree, tar_file)

class DirFileWriter(object):

    def __init__(self, fileobj, mode):

        # We may be asked to 'append'.  If so, fileobj already has a path.

        # So we copy the existing tree, and overwrite afterward.

        fileobj.seek(0)

        existing = fileobj.read()

        fileobj.seek(0)

        path = tempfile.mkdtemp(dir=os.getcwd())

        if existing != '':

            # copytree requires the directory not to exist

            os.rmdir(path)

            copytree(existing, path)

        fileobj.write(path)

        self.root = path

    def add(self, path):

        target_path = os.path.join(self.root, path)

        parent = osutils.dirname(target_path)

        if not os.path.exists(parent):

            os.makedirs(parent)

        kind = osutils.file_kind(path)

        if kind == 'file':

            copy2(path, target_path)

        if kind == 'directory':

            os.mkdir(target_path)

    def close(self):

        pass

class TestImport(TestCaseInTempDir):

    def make_tar(self, mode='w'):

        def maker(fileobj):

            return tarfile.open('project-0.1.tar', mode, fileobj)

        return self.make_archive(maker)

    def make_archive(self, maker, subdir=True):

        result = StringIO()

        archive_file = maker(result)

        try:

            os.mkdir('project-0.1')

            if subdir:

                prefix='project-0.1/'

                archive_file.add('project-0.1')

            else:

                prefix=''

                os.chdir('project-0.1')

            os.mkdir(prefix + 'junk')

            archive_file.add(prefix + 'junk')

            f = file(prefix + 'README', 'wb')

            f.write('What?')

            f.close()

            archive_file.add(prefix + 'README')

            f = file(prefix + 'FEEDME', 'wb')

            f.write('Hungry!!')

            f.close()

            archive_file.add(prefix + 'FEEDME')

            archive_file.close()

        finally:

            if not subdir:

                os.chdir('..')

            rmtree('project-0.1')

        result.seek(0)

        return result

    def make_archive2(self, builder, subdir):

        result = StringIO()

        archive_file = builder(result)

        os.mkdir('project-0.2')

        try:

            if subdir:

                prefix='project-0.2/'

                archive_file.add('project-0.2')

            else:

                prefix=''

                os.chdir('project-0.2')

            os.mkdir(prefix + 'junk')

            archive_file.add(prefix + 'junk')

            f = file(prefix + 'README', 'wb')

            f.write('Now?')

            f.close()

            archive_file.add(prefix + 'README')

            archive_file.close()

            archive_file = builder(result, 'a')

            archive_file.add(prefix + 'README')

            archive_file.close()

        finally:

            if not subdir:

                os.chdir('..')

        result.seek(0)

        return result

    def make_messed_tar(self):

        result = StringIO()

        tar_file = tarfile.open('project-0.1.tar', 'w', result)

        os.mkdir('project-0.1')

        tar_file.add('project-0.1')

        os.mkdir('project-0.2')

        tar_file.add('project-0.2')

        f = file('project-0.1/README', 'wb')

        f.write('What?')

        f.close()

        tar_file.add('project-0.1/README')

        tar_file.close()

        rmtree('project-0.1')

        result.seek(0)

        return result

    def make_zip(self):

        def maker(fileobj):

            return ZipFileWrapper(fileobj, 'w')

        return self.make_archive(maker)

    def test_top_directory(self):

        self.assertEqual(top_directory('ab/b/c'), 'ab')

        self.assertEqual(top_directory('/etc'), '/')

    def test_common_directory(self):

        self.assertEqual(common_directory(['ab/c/d', 'ab/c/e']), 'ab')

        self.assertIs(common_directory(['ab/c/d', 'ac/c/e']), None)

        self.assertIs(None, common_directory(['FEEDME']))

    def test_untar(self):

        def builder(fileobj, mode='w'):

            return tarfile.open('project-0.1.tar', mode, fileobj)

        self.archive_test(builder, import_tar)

    def test_broken_tar(self):

        def builder(fileobj, mode='w'):

            return tarfile.open('project-0.1.tar', mode, fileobj)

        self.archive_test(builder, import_tar_broken, subdir=True)

    def test_unzip(self):

        def builder(fileobj, mode='w'):

            return ZipFileWrapper(fileobj, mode)

        self.archive_test(builder, import_zip)

    def test_copydir_nosub(self):

        def builder(fileobj, mode='w'):

            return DirFileWriter(fileobj, mode)

        # It would be bogus to test with the result in a subdirectory,

        # because for directories, the input root is always the output root.

        self.archive_test(builder, import_dir)

    def archive_test(self, builder, importer, subdir=False):

        archive_file = self.make_archive(builder, subdir)

        tree = BzrDir.create_standalone_workingtree('tree')

        tree.lock_write()

        try:

            importer(tree, archive_file)

            self.assertTrue(tree.path2id('README') is not None)

            self.assertTrue(tree.path2id('FEEDME') is not None)

            self.assertTrue(os.path.isfile(tree.abspath('README')))

            self.assertEqual(tree.inventory[tree.path2id('README')].kind,

                'file')

            self.assertEqual(tree.inventory[tree.path2id('FEEDME')].kind,

                'file')

            f = file(tree.abspath('junk/food'), 'wb')

            f.write('I like food\n')

            f.close()

            archive_file = self.make_archive2(builder, subdir)

            importer(tree, archive_file)

            self.assertTrue(tree.path2id('README') is not None)

            self.assertTrue(not os.path.exists(tree.abspath('FEEDME')))

        finally:

            tree.unlock()

    def test_untar2(self):

        tar_file = self.make_messed_tar()

        tree = BzrDir.create_standalone_workingtree('tree')

        import_tar(tree, tar_file)

        self.assertTrue(tree.path2id('project-0.1/README') is not None) 

    def test_untar_gzip(self):

        tar_file = self.make_tar(mode='w:gz')

        tree = BzrDir.create_standalone_workingtree('tree')

        import_tar(tree, tar_file)

        self.assertTrue(tree.path2id('README') is not None) 

def test_suite():

    return makeSuite(TestImport)