michael@0: #!/usr/bin/env python
michael@0: 
michael@0: import mozfile
michael@0: import os
michael@0: import shutil
michael@0: import tarfile
michael@0: import tempfile
michael@0: import stubs
michael@0: import unittest
michael@0: import zipfile
michael@0: 
michael@0: 
michael@0: class TestExtract(unittest.TestCase):
michael@0:     """test extracting archives"""
michael@0: 
michael@0:     def ensure_directory_contents(self, directory):
michael@0:         """ensure the directory contents match"""
michael@0:         for f in stubs.files:
michael@0:             path = os.path.join(directory, *f)
michael@0:             exists = os.path.exists(path)
michael@0:             if not exists:
michael@0:                 print "%s does not exist" % (os.path.join(f))
michael@0:             self.assertTrue(exists)
michael@0:             if exists:
michael@0:                 contents = file(path).read().strip()
michael@0:                 self.assertTrue(contents == f[-1])
michael@0: 
michael@0:     def test_extract_zipfile(self):
michael@0:         """test extracting a zipfile"""
michael@0:         _zipfile = self.create_zip()
michael@0:         self.assertTrue(os.path.exists(_zipfile))
michael@0:         try:
michael@0:             dest = tempfile.mkdtemp()
michael@0:             try:
michael@0:                 mozfile.extract_zip(_zipfile, dest)
michael@0:                 self.ensure_directory_contents(dest)
michael@0:             finally:
michael@0:                 shutil.rmtree(dest)
michael@0:         finally:
michael@0:             os.remove(_zipfile)
michael@0: 
michael@0:     def test_extract_tarball(self):
michael@0:         """test extracting a tarball"""
michael@0:         tarball = self.create_tarball()
michael@0:         self.assertTrue(os.path.exists(tarball))
michael@0:         try:
michael@0:             dest = tempfile.mkdtemp()
michael@0:             try:
michael@0:                 mozfile.extract_tarball(tarball, dest)
michael@0:                 self.ensure_directory_contents(dest)
michael@0:             finally:
michael@0:                 shutil.rmtree(dest)
michael@0:         finally:
michael@0:             os.remove(tarball)
michael@0: 
michael@0:     def test_extract(self):
michael@0:         """test the generalized extract function"""
michael@0: 
michael@0:         # test extracting a tarball
michael@0:         tarball = self.create_tarball()
michael@0:         self.assertTrue(os.path.exists(tarball))
michael@0:         try:
michael@0:             dest = tempfile.mkdtemp()
michael@0:             try:
michael@0:                 mozfile.extract(tarball, dest)
michael@0:                 self.ensure_directory_contents(dest)
michael@0:             finally:
michael@0:                 shutil.rmtree(dest)
michael@0:         finally:
michael@0:             os.remove(tarball)
michael@0: 
michael@0:         # test extracting a zipfile
michael@0:         _zipfile = self.create_zip()
michael@0:         self.assertTrue(os.path.exists(_zipfile))
michael@0:         try:
michael@0:             dest = tempfile.mkdtemp()
michael@0:             try:
michael@0:                 mozfile.extract_zip(_zipfile, dest)
michael@0:                 self.ensure_directory_contents(dest)
michael@0:             finally:
michael@0:                 shutil.rmtree(dest)
michael@0:         finally:
michael@0:             os.remove(_zipfile)
michael@0: 
michael@0:         # test extracting some non-archive; this should fail
michael@0:         fd, filename = tempfile.mkstemp()
michael@0:         os.write(fd, 'This is not a zipfile or tarball')
michael@0:         os.close(fd)
michael@0:         exception = None
michael@0:         try:
michael@0:             dest = tempfile.mkdtemp()
michael@0:             mozfile.extract(filename, dest)
michael@0:         except Exception, exception:
michael@0:             pass
michael@0:         finally:
michael@0:             os.remove(filename)
michael@0:             os.rmdir(dest)
michael@0:         self.assertTrue(isinstance(exception, Exception))
michael@0: 
michael@0:     ### utility functions
michael@0: 
michael@0:     def create_tarball(self):
michael@0:         """create a stub tarball for testing"""
michael@0:         tempdir = stubs.create_stub()
michael@0:         filename = tempfile.mktemp(suffix='.tar')
michael@0:         archive = tarfile.TarFile(filename, mode='w')
michael@0:         try:
michael@0:             for path in stubs.files:
michael@0:                 archive.add(os.path.join(tempdir, *path), arcname=os.path.join(*path))
michael@0:         except:
michael@0:             os.remove(archive)
michael@0:             raise
michael@0:         finally:
michael@0:             shutil.rmtree(tempdir)
michael@0:         archive.close()
michael@0:         return filename
michael@0: 
michael@0:     def create_zip(self):
michael@0:         """create a stub zipfile for testing"""
michael@0: 
michael@0:         tempdir = stubs.create_stub()
michael@0:         filename = tempfile.mktemp(suffix='.zip')
michael@0:         archive = zipfile.ZipFile(filename, mode='w')
michael@0:         try:
michael@0:             for path in stubs.files:
michael@0:                 archive.write(os.path.join(tempdir, *path), arcname=os.path.join(*path))
michael@0:         except:
michael@0:             os.remove(filename)
michael@0:             raise
michael@0:         finally:
michael@0:             shutil.rmtree(tempdir)
michael@0:         archive.close()
michael@0:         return filename