#!/usr/bin/env python

import mozfile
import os
import shutil
import tarfile
import tempfile
import stubs
import unittest
import zipfile


class TestExtract(unittest.TestCase):
    """test extracting archives"""

    def ensure_directory_contents(self, directory):
        """ensure the directory contents match"""
        for f in stubs.files:
            path = os.path.join(directory, *f)
            exists = os.path.exists(path)
            if not exists:
                print "%s does not exist" % (os.path.join(f))
            self.assertTrue(exists)
            if exists:
                contents = file(path).read().strip()
                self.assertTrue(contents == f[-1])

    def test_extract_zipfile(self):
        """test extracting a zipfile"""
        _zipfile = self.create_zip()
        self.assertTrue(os.path.exists(_zipfile))
        try:
            dest = tempfile.mkdtemp()
            try:
                mozfile.extract_zip(_zipfile, dest)
                self.ensure_directory_contents(dest)
            finally:
                shutil.rmtree(dest)
        finally:
            os.remove(_zipfile)

    def test_extract_tarball(self):
        """test extracting a tarball"""
        tarball = self.create_tarball()
        self.assertTrue(os.path.exists(tarball))
        try:
            dest = tempfile.mkdtemp()
            try:
                mozfile.extract_tarball(tarball, dest)
                self.ensure_directory_contents(dest)
            finally:
                shutil.rmtree(dest)
        finally:
            os.remove(tarball)

    def test_extract(self):
        """test the generalized extract function"""

        # test extracting a tarball
        tarball = self.create_tarball()
        self.assertTrue(os.path.exists(tarball))
        try:
            dest = tempfile.mkdtemp()
            try:
                mozfile.extract(tarball, dest)
                self.ensure_directory_contents(dest)
            finally:
                shutil.rmtree(dest)
        finally:
            os.remove(tarball)

        # test extracting a zipfile
        _zipfile = self.create_zip()
        self.assertTrue(os.path.exists(_zipfile))
        try:
            dest = tempfile.mkdtemp()
            try:
                mozfile.extract_zip(_zipfile, dest)
                self.ensure_directory_contents(dest)
            finally:
                shutil.rmtree(dest)
        finally:
            os.remove(_zipfile)

        # test extracting some non-archive; this should fail
        fd, filename = tempfile.mkstemp()
        os.write(fd, 'This is not a zipfile or tarball')
        os.close(fd)
        exception = None
        try:
            dest = tempfile.mkdtemp()
            mozfile.extract(filename, dest)
        except Exception, exception:
            pass
        finally:
            os.remove(filename)
            os.rmdir(dest)
        self.assertTrue(isinstance(exception, Exception))

    ### utility functions

    def create_tarball(self):
        """create a stub tarball for testing"""
        tempdir = stubs.create_stub()
        filename = tempfile.mktemp(suffix='.tar')
        archive = tarfile.TarFile(filename, mode='w')
        try:
            for path in stubs.files:
                archive.add(os.path.join(tempdir, *path), arcname=os.path.join(*path))
        except:
            os.remove(archive)
            raise
        finally:
            shutil.rmtree(tempdir)
        archive.close()
        return filename

    def create_zip(self):
        """create a stub zipfile for testing"""

        tempdir = stubs.create_stub()
        filename = tempfile.mktemp(suffix='.zip')
        archive = zipfile.ZipFile(filename, mode='w')
        try:
            for path in stubs.files:
                archive.write(os.path.join(tempdir, *path), arcname=os.path.join(*path))
        except:
            os.remove(filename)
            raise
        finally:
            shutil.rmtree(tempdir)
        archive.close()
        return filename