1#!/usr/bin/python2.4
2#
3# Copyright (C) 2008 Google Inc.
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#      http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16#
17
18"""Tests for divide_and_compress.py.
19
20TODO(jmatt): Add tests for module methods.
21"""
22
23__author__ = 'jmatt@google.com (Justin Mattson)'
24
25import os
26import stat
27import unittest
28import zipfile
29
30import divide_and_compress
31import mox
32
33
34class BagOfParts(object):
35  """Just a generic class that I can use to assign random attributes to."""
36
37  def NoOp(self):
38    x = 1
39
40
41class ValidAndRemoveTests(unittest.TestCase):
42  """Test the ArchiveIsValid and RemoveLastFile methods."""
43
44  def setUp(self):
45    """Prepare the test.
46
47    Construct some mock objects for use with the tests.
48    """
49    self.my_mox = mox.Mox()
50    file1 = BagOfParts()
51    file1.filename = 'file1.txt'
52    file1.contents = 'This is a test file'
53    file2 = BagOfParts()
54    file2.filename = 'file2.txt'
55    file2.contents = ('akdjfk;djsf;kljdslkfjslkdfjlsfjkdvn;kn;2389rtu4i'
56                      'tn;ghf8:89H*hp748FJw80fu9WJFpwf39pujens;fihkhjfk'
57                      'sdjfljkgsc n;iself')
58    self.files = {'file1': file1, 'file2': file2}
59
60  def tearDown(self):
61    """Remove any stubs we've created."""
62    self.my_mox.UnsetStubs()
63
64  def testArchiveIsValid(self):
65    """Test the DirectoryZipper.ArchiveIsValid method.
66
67    Run two tests, one that we expect to pass and one that we expect to fail
68    """
69    test_file_size = 1056730
70    self.my_mox.StubOutWithMock(os, 'stat')
71    os.stat('/foo/0.zip').AndReturn([test_file_size])
72    self.my_mox.StubOutWithMock(stat, 'ST_SIZE')
73    stat.ST_SIZE = 0
74    os.stat('/baz/0.zip').AndReturn([test_file_size])
75    mox.Replay(os.stat)
76    test_target = divide_and_compress.DirectoryZipper('/foo/', 'bar',
77                                                      test_file_size - 1, True)
78
79    self.assertEqual(False, test_target.ArchiveIsValid(),
80                     msg=('ERROR: Test failed, ArchiveIsValid should have '
81                          'returned false, but returned true'))
82
83    test_target = divide_and_compress.DirectoryZipper('/baz/', 'bar',
84                                                      test_file_size + 1, True)
85    self.assertEqual(True, test_target.ArchiveIsValid(),
86                     msg=('ERROR: Test failed, ArchiveIsValid should have'
87                          ' returned true, but returned false'))
88
89  def testRemoveLastFile(self):
90    """Test DirectoryZipper.RemoveLastFile method.
91
92    Construct a ZipInfo mock object with two records, verify that write is
93    only called once on the new ZipFile object.
94    """
95    source = self.CreateZipSource()
96    dest = self.CreateZipDestination()
97    source_path = ''.join([os.getcwd(), '/0-old.zip'])
98    dest_path = ''.join([os.getcwd(), '/0.zip'])
99    test_target = divide_and_compress.DirectoryZipper(
100        ''.join([os.getcwd(), '/']), 'dummy', 1024*1024, True)
101    self.my_mox.StubOutWithMock(test_target, 'OpenZipFileAtPath')
102    test_target.OpenZipFileAtPath(source_path, mode='r').AndReturn(source)
103    test_target.OpenZipFileAtPath(dest_path,
104                                  compress=zipfile.ZIP_DEFLATED,
105                                  mode='w').AndReturn(dest)
106    self.my_mox.StubOutWithMock(os, 'rename')
107    os.rename(dest_path, source_path)
108    self.my_mox.StubOutWithMock(os, 'unlink')
109    os.unlink(source_path)
110
111    self.my_mox.ReplayAll()
112    test_target.RemoveLastFile()
113    self.my_mox.VerifyAll()
114
115  def CreateZipSource(self):
116    """Create a mock zip sourec object.
117
118    Read should only be called once, because the second file is the one
119    being removed.
120
121    Returns:
122      A configured mocked
123    """
124
125    source_zip = self.my_mox.CreateMock(zipfile.ZipFile)
126    source_zip.infolist().AndReturn([self.files['file1'], self.files['file1']])
127    source_zip.infolist().AndReturn([self.files['file1'], self.files['file1']])
128    source_zip.read(self.files['file1'].filename).AndReturn(
129        self.files['file1'].contents)
130    source_zip.close()
131    return source_zip
132
133  def CreateZipDestination(self):
134    """Create mock destination zip.
135
136    Write should only be called once, because there are two files in the
137    source zip and we expect the second to be removed.
138
139    Returns:
140      A configured mocked
141    """
142
143    dest_zip = mox.MockObject(zipfile.ZipFile)
144    dest_zip.writestr(self.files['file1'].filename,
145                      self.files['file1'].contents)
146    dest_zip.close()
147    return dest_zip
148
149
150class FixArchiveTests(unittest.TestCase):
151  """Tests for the DirectoryZipper.FixArchive method."""
152
153  def setUp(self):
154    """Create a mock file object."""
155    self.my_mox = mox.Mox()
156    self.file1 = BagOfParts()
157    self.file1.filename = 'file1.txt'
158    self.file1.contents = 'This is a test file'
159
160  def tearDown(self):
161    """Unset any mocks that we've created."""
162    self.my_mox.UnsetStubs()
163
164  def _InitMultiFileData(self):
165    """Create an array of mock file objects.
166
167    Create three mock file objects that we can use for testing.
168    """
169    self.multi_file_dir = []
170
171    file1 = BagOfParts()
172    file1.filename = 'file1.txt'
173    file1.contents = 'kjaskl;jkdjfkja;kjsnbvjnvnbuewklriujalvjsd'
174    self.multi_file_dir.append(file1)
175
176    file2 = BagOfParts()
177    file2.filename = 'file2.txt'
178    file2.contents = ('He entered the room and there in the center, it was.'
179                      ' Looking upon the thing, suddenly he could not remember'
180                      ' whether he had actually seen it before or whether'
181                      ' his memory of it was merely the effect of something'
182                      ' so often being imagined that it had long since become '
183                      ' manifest in his mind.')
184    self.multi_file_dir.append(file2)
185
186    file3 = BagOfParts()
187    file3.filename = 'file3.txt'
188    file3.contents = 'Whoa, what is \'file2.txt\' all about?'
189    self.multi_file_dir.append(file3)
190
191  def testSingleFileArchive(self):
192    """Test behavior of FixArchive when the archive has a single member.
193
194    We expect that when this method is called with an archive that has a
195    single member that it will return False and unlink the archive.
196    """
197    test_target = divide_and_compress.DirectoryZipper(
198        ''.join([os.getcwd(), '/']), 'dummy', 1024*1024, True)
199    self.my_mox.StubOutWithMock(test_target, 'OpenZipFileAtPath')
200    test_target.OpenZipFileAtPath(
201        ''.join([os.getcwd(), '/0.zip']), mode='r').AndReturn(
202            self.CreateSingleFileMock())
203    self.my_mox.StubOutWithMock(os, 'unlink')
204    os.unlink(''.join([os.getcwd(), '/0.zip']))
205    self.my_mox.ReplayAll()
206    self.assertEqual(False, test_target.FixArchive('SIZE'))
207    self.my_mox.VerifyAll()
208
209  def CreateSingleFileMock(self):
210    """Create a mock ZipFile object for testSingleFileArchive.
211
212    We just need it to return a single member infolist twice
213
214    Returns:
215      A configured mock object
216    """
217    mock_zip = self.my_mox.CreateMock(zipfile.ZipFile)
218    mock_zip.infolist().AndReturn([self.file1])
219    mock_zip.infolist().AndReturn([self.file1])
220    mock_zip.close()
221    return mock_zip
222
223  def testMultiFileArchive(self):
224    """Test behavior of DirectoryZipper.FixArchive with a multi-file archive.
225
226    We expect that FixArchive will rename the old archive, adding '-old' before
227    '.zip', read all the members except the last one of '-old' into a new
228    archive with the same name as the original, and then unlink the '-old' copy
229    """
230    test_target = divide_and_compress.DirectoryZipper(
231        ''.join([os.getcwd(), '/']), 'dummy', 1024*1024, True)
232    self.my_mox.StubOutWithMock(test_target, 'OpenZipFileAtPath')
233    test_target.OpenZipFileAtPath(
234        ''.join([os.getcwd(), '/0.zip']), mode='r').AndReturn(
235            self.CreateMultiFileMock())
236    self.my_mox.StubOutWithMock(test_target, 'RemoveLastFile')
237    test_target.RemoveLastFile(''.join([os.getcwd(), '/0.zip']))
238    self.my_mox.StubOutWithMock(os, 'stat')
239    os.stat(''.join([os.getcwd(), '/0.zip'])).AndReturn([49302])
240    self.my_mox.StubOutWithMock(stat, 'ST_SIZE')
241    stat.ST_SIZE = 0
242    self.my_mox.ReplayAll()
243    self.assertEqual(True, test_target.FixArchive('SIZE'))
244    self.my_mox.VerifyAll()
245
246  def CreateMultiFileMock(self):
247    """Create mock ZipFile object for use with testMultiFileArchive.
248
249    The mock just needs to return the infolist mock that is prepared in
250    InitMultiFileData()
251
252    Returns:
253      A configured mock object
254    """
255    self._InitMultiFileData()
256    mock_zip = self.my_mox.CreateMock(zipfile.ZipFile)
257    mock_zip.infolist().AndReturn(self.multi_file_dir)
258    mock_zip.close()
259    return mock_zip
260
261
262class AddFileToArchiveTest(unittest.TestCase):
263  """Test behavior of method to add a file to an archive."""
264
265  def setUp(self):
266    """Setup the arguments for the DirectoryZipper object."""
267    self.my_mox = mox.Mox()
268    self.output_dir = '%s/' % os.getcwd()
269    self.file_to_add = 'file.txt'
270    self.input_dir = '/foo/bar/baz/'
271
272  def tearDown(self):
273    self.my_mox.UnsetStubs()
274
275  def testAddFileToArchive(self):
276    """Test the DirectoryZipper.AddFileToArchive method.
277
278    We are testing a pretty trivial method, we just expect it to look at the
279    file its adding, so that it possible can through out a warning.
280    """
281    test_target = divide_and_compress.DirectoryZipper(self.output_dir,
282                                                      self.input_dir,
283                                                      1024*1024, True)
284    self.my_mox.StubOutWithMock(test_target, 'OpenZipFileAtPath')
285    archive_mock = self.CreateArchiveMock()
286    test_target.OpenZipFileAtPath(
287        ''.join([self.output_dir, '0.zip']),
288        compress=zipfile.ZIP_DEFLATED).AndReturn(archive_mock)
289    self.StubOutOsModule()
290    self.my_mox.ReplayAll()
291    test_target.AddFileToArchive(''.join([self.input_dir, self.file_to_add]),
292                                 zipfile.ZIP_DEFLATED)
293    self.my_mox.VerifyAll()
294
295  def StubOutOsModule(self):
296    """Create a mock for the os.path and os.stat objects.
297
298    Create a stub that will return the type (file or directory) and size of the
299    object that is to be added.
300    """
301    self.my_mox.StubOutWithMock(os.path, 'isfile')
302    os.path.isfile(''.join([self.input_dir, self.file_to_add])).AndReturn(True)
303    self.my_mox.StubOutWithMock(os, 'stat')
304    os.stat(''.join([self.input_dir, self.file_to_add])).AndReturn([39480])
305    self.my_mox.StubOutWithMock(stat, 'ST_SIZE')
306    stat.ST_SIZE = 0
307
308  def CreateArchiveMock(self):
309    """Create a mock ZipFile for use with testAddFileToArchive.
310
311    Just verify that write is called with the file we expect and that the
312    archive is closed after the file addition
313
314    Returns:
315      A configured mock object
316    """
317    archive_mock = self.my_mox.CreateMock(zipfile.ZipFile)
318    archive_mock.write(''.join([self.input_dir, self.file_to_add]),
319                       self.file_to_add)
320    archive_mock.close()
321    return archive_mock
322
323
324class CompressDirectoryTest(unittest.TestCase):
325  """Test the master method of the class.
326
327  Testing with the following directory structure.
328  /dir1/
329  /dir1/file1.txt
330  /dir1/file2.txt
331  /dir1/dir2/
332  /dir1/dir2/dir3/
333  /dir1/dir2/dir4/
334  /dir1/dir2/dir4/file3.txt
335  /dir1/dir5/
336  /dir1/dir5/file4.txt
337  /dir1/dir5/file5.txt
338  /dir1/dir5/file6.txt
339  /dir1/dir5/file7.txt
340  /dir1/dir6/
341  /dir1/dir6/file8.txt
342
343  file1.txt., file2.txt, file3.txt should be in 0.zip
344  file4.txt should be in 1.zip
345  file5.txt, file6.txt should be in 2.zip
346  file7.txt will not be stored since it will be too large compressed
347  file8.txt should b in 3.zip
348  """
349
350  def setUp(self):
351    """Setup all the mocks for this test."""
352    self.my_mox = mox.Mox()
353
354    self.base_dir = '/dir1'
355    self.output_path = '/out_dir/'
356    self.test_target = divide_and_compress.DirectoryZipper(
357        self.output_path, self.base_dir, 1024*1024, True)
358
359    self.InitArgLists()
360    self.InitOsDotPath()
361    self.InitArchiveIsValid()
362    self.InitWriteIndexRecord()
363    self.InitAddFileToArchive()
364
365  def tearDown(self):
366    self.my_mox.UnsetStubs()
367
368  def testCompressDirectory(self):
369    """Test the DirectoryZipper.CompressDirectory method."""
370    self.my_mox.ReplayAll()
371    for arguments in self.argument_lists:
372      self.test_target.CompressDirectory(None, arguments[0], arguments[1])
373    self.my_mox.VerifyAll()
374
375  def InitAddFileToArchive(self):
376    """Setup mock for DirectoryZipper.AddFileToArchive.
377
378    Make sure that the files are added in the order we expect.
379    """
380    self.my_mox.StubOutWithMock(self.test_target, 'AddFileToArchive')
381    self.test_target.AddFileToArchive('/dir1/file1.txt', zipfile.ZIP_DEFLATED)
382    self.test_target.AddFileToArchive('/dir1/file2.txt', zipfile.ZIP_DEFLATED)
383    self.test_target.AddFileToArchive('/dir1/dir2/dir4/file3.txt',
384                                      zipfile.ZIP_DEFLATED)
385    self.test_target.AddFileToArchive('/dir1/dir5/file4.txt',
386                                      zipfile.ZIP_DEFLATED)
387    self.test_target.AddFileToArchive('/dir1/dir5/file4.txt',
388                                      zipfile.ZIP_DEFLATED)
389    self.test_target.AddFileToArchive('/dir1/dir5/file5.txt',
390                                      zipfile.ZIP_DEFLATED)
391    self.test_target.AddFileToArchive('/dir1/dir5/file5.txt',
392                                      zipfile.ZIP_DEFLATED)
393    self.test_target.AddFileToArchive('/dir1/dir5/file6.txt',
394                                      zipfile.ZIP_DEFLATED)
395    self.test_target.AddFileToArchive('/dir1/dir5/file7.txt',
396                                      zipfile.ZIP_DEFLATED)
397    self.test_target.AddFileToArchive('/dir1/dir5/file7.txt',
398                                      zipfile.ZIP_DEFLATED)
399    self.test_target.AddFileToArchive('/dir1/dir6/file8.txt',
400                                      zipfile.ZIP_DEFLATED)
401
402  def InitWriteIndexRecord(self):
403    """Setup mock for DirectoryZipper.WriteIndexRecord."""
404    self.my_mox.StubOutWithMock(self.test_target, 'WriteIndexRecord')
405
406    # we are trying to compress 8 files, but we should only attempt to
407    # write an index record 7 times, because one file is too large to be stored
408    self.test_target.WriteIndexRecord().AndReturn(True)
409    self.test_target.WriteIndexRecord().AndReturn(False)
410    self.test_target.WriteIndexRecord().AndReturn(False)
411    self.test_target.WriteIndexRecord().AndReturn(True)
412    self.test_target.WriteIndexRecord().AndReturn(True)
413    self.test_target.WriteIndexRecord().AndReturn(False)
414    self.test_target.WriteIndexRecord().AndReturn(True)
415
416  def InitArchiveIsValid(self):
417    """Mock out DirectoryZipper.ArchiveIsValid and DirectoryZipper.FixArchive.
418
419    Mock these methods out such that file1, file2, and file3 go into one
420    archive. file4 then goes into the next archive, file5 and file6 in the
421    next, file 7 should appear too large to compress into an archive, and
422    file8 goes into the final archive
423    """
424    self.my_mox.StubOutWithMock(self.test_target, 'ArchiveIsValid')
425    self.my_mox.StubOutWithMock(self.test_target, 'FixArchive')
426    self.test_target.ArchiveIsValid().AndReturn(True)
427    self.test_target.ArchiveIsValid().AndReturn(True)
428    self.test_target.ArchiveIsValid().AndReturn(True)
429
430    # should be file4.txt
431    self.test_target.ArchiveIsValid().AndReturn(False)
432    self.test_target.FixArchive('SIZE').AndReturn(True)
433    self.test_target.ArchiveIsValid().AndReturn(True)
434
435    # should be file5.txt
436    self.test_target.ArchiveIsValid().AndReturn(False)
437    self.test_target.FixArchive('SIZE').AndReturn(True)
438    self.test_target.ArchiveIsValid().AndReturn(True)
439    self.test_target.ArchiveIsValid().AndReturn(True)
440
441    # should be file7.txt
442    self.test_target.ArchiveIsValid().AndReturn(False)
443    self.test_target.FixArchive('SIZE').AndReturn(True)
444    self.test_target.ArchiveIsValid().AndReturn(False)
445    self.test_target.FixArchive('SIZE').AndReturn(False)
446    self.test_target.ArchiveIsValid().AndReturn(True)
447
448  def InitOsDotPath(self):
449    """Mock out os.path.isfile.
450
451    Mock this out so the things we want to appear as files appear as files and
452    the things we want to appear as directories appear as directories. Also
453    make sure that the order of file visits is as we expect (which is why
454    InAnyOrder isn't used here).
455    """
456    self.my_mox.StubOutWithMock(os.path, 'isfile')
457    os.path.isfile('/dir1/dir2').AndReturn(False)
458    os.path.isfile('/dir1/dir5').AndReturn(False)
459    os.path.isfile('/dir1/dir6').AndReturn(False)
460    os.path.isfile('/dir1/file1.txt').AndReturn(True)
461    os.path.isfile('/dir1/file2.txt').AndReturn(True)
462    os.path.isfile('/dir1/dir2/dir3').AndReturn(False)
463    os.path.isfile('/dir1/dir2/dir4').AndReturn(False)
464    os.path.isfile('/dir1/dir2/dir4/file3.txt').AndReturn(True)
465    os.path.isfile('/dir1/dir5/file4.txt').AndReturn(True)
466    os.path.isfile('/dir1/dir5/file4.txt').AndReturn(True)
467    os.path.isfile('/dir1/dir5/file5.txt').AndReturn(True)
468    os.path.isfile('/dir1/dir5/file5.txt').AndReturn(True)
469    os.path.isfile('/dir1/dir5/file6.txt').AndReturn(True)
470    os.path.isfile('/dir1/dir5/file7.txt').AndReturn(True)
471    os.path.isfile('/dir1/dir5/file7.txt').AndReturn(True)
472    os.path.isfile('/dir1/dir6/file8.txt').AndReturn(True)
473
474  def InitArgLists(self):
475    """Create the directory path => directory contents mappings."""
476    self.argument_lists = []
477    self.argument_lists.append(['/dir1',
478                                ['file1.txt', 'file2.txt', 'dir2', 'dir5',
479                                 'dir6']])
480    self.argument_lists.append(['/dir1/dir2', ['dir3', 'dir4']])
481    self.argument_lists.append(['/dir1/dir2/dir3', []])
482    self.argument_lists.append(['/dir1/dir2/dir4', ['file3.txt']])
483    self.argument_lists.append(['/dir1/dir5',
484                                ['file4.txt', 'file5.txt', 'file6.txt',
485                                 'file7.txt']])
486    self.argument_lists.append(['/dir1/dir6', ['file8.txt']])
487
488if __name__ == '__main__':
489  unittest.main()
490