michael@0: # This Source Code Form is subject to the terms of the Mozilla Public
michael@0: # License, v. 2.0. If a copy of the MPL was not distributed with this
michael@0: # file, You can obtain one at http://mozilla.org/MPL/2.0/.
michael@0: 
michael@0: 
michael@0: from random import randint
michael@0: from zipfile import ZipFile
michael@0: import os
michael@0: import shutil
michael@0: 
michael@0: 
michael@0: def gen_binary_file(path, size):
michael@0:     with open(path, 'wb') as f:
michael@0:         for i in xrange(size):
michael@0:             byte = '%c' % randint(0, 255)
michael@0:             f.write(byte)
michael@0: 
michael@0: 
michael@0: def gen_zip(path, files, stripped_prefix=''):
michael@0:     with ZipFile(path, 'w') as z:
michael@0:         for f in files:
michael@0:             new_name = f.replace(stripped_prefix, '')
michael@0:             z.write(f, new_name)
michael@0: 
michael@0: 
michael@0: def mkdir(path, *args):
michael@0:     try:
michael@0:         os.mkdir(path, *args)
michael@0:     except OSError:
michael@0:         pass
michael@0: 
michael@0: 
michael@0: def gen_folder_structure():
michael@0:     root = 'test-files'
michael@0:     prefix = os.path.join(root, 'push2')
michael@0:     mkdir(prefix)
michael@0: 
michael@0:     gen_binary_file(os.path.join(prefix, 'file4.bin'), 59036)
michael@0:     mkdir(os.path.join(prefix, 'sub1'))
michael@0:     shutil.copyfile(os.path.join(root, 'mytext.txt'),
michael@0:                     os.path.join(prefix, 'sub1', 'file1.txt'))
michael@0:     mkdir(os.path.join(prefix, 'sub1', 'sub1.1'))
michael@0:     shutil.copyfile(os.path.join(root, 'mytext.txt'),
michael@0:                     os.path.join(prefix, 'sub1', 'sub1.1', 'file2.txt'))
michael@0:     mkdir(os.path.join(prefix, 'sub2'))
michael@0:     shutil.copyfile(os.path.join(root, 'mytext.txt'),
michael@0:                     os.path.join(prefix, 'sub2', 'file3.txt'))
michael@0: 
michael@0: 
michael@0: def gen_test_files():
michael@0:     gen_folder_structure()
michael@0:     flist = [
michael@0:         os.path.join('test-files', 'push2'),
michael@0:         os.path.join('test-files', 'push2', 'file4.bin'),
michael@0:         os.path.join('test-files', 'push2', 'sub1'),
michael@0:         os.path.join('test-files', 'push2', 'sub1', 'file1.txt'),
michael@0:         os.path.join('test-files', 'push2', 'sub1', 'sub1.1'),
michael@0:         os.path.join('test-files', 'push2', 'sub1', 'sub1.1', 'file2.txt'),
michael@0:         os.path.join('test-files', 'push2', 'sub2'),
michael@0:         os.path.join('test-files', 'push2', 'sub2', 'file3.txt')
michael@0:     ]
michael@0:     gen_zip(os.path.join('test-files', 'mybinary.zip'),
michael@0:             flist, stripped_prefix=('test-files' + os.path.sep))
michael@0:     gen_zip(os.path.join('test-files', 'mytext.zip'),
michael@0:             [os.path.join('test-files', 'mytext.txt')])
michael@0: 
michael@0: 
michael@0: def clean_test_files():
michael@0:     ds = [os.path.join('test-files', d) for d in ('push1', 'push2')]
michael@0:     for d in ds:
michael@0:         try:
michael@0:             shutil.rmtree(d)
michael@0:         except OSError:
michael@0:             pass
michael@0: 
michael@0:     fs = [os.path.join('test-files', f) for f in ('mybinary.zip', 'mytext.zip')]
michael@0:     for f in fs:
michael@0:         try:
michael@0:             os.remove(f)
michael@0:         except OSError:
michael@0:             pass
michael@0: 
michael@0: 
michael@0: if __name__ == '__main__':
michael@0:     gen_test_files()