blob: d9e6e0750b6cbee407e73a9a0ba260a876267322 [file] [log] [blame]
#!/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