build/zip_helpers_unittest.py

   1 #!/usr/bin/env python3
   2 # Copyright 2023 The Chromium Authors
   3 # Use of this source code is governed by a BSD-style license that can be
   4 # found in the LICENSE file.
   5
   6 import os
   7 import pathlib
   8 import shutil
   9 import sys
  10 import tempfile
  11 import unittest
  12 import zipfile
  13
  14 import zip_helpers
  15
  16
  17 def _make_test_zips(tmp_dir, create_conflct=False):
  18   zip1 = os.path.join(tmp_dir, 'A.zip')
  19   zip2 = os.path.join(tmp_dir, 'B.zip')
  20   with zipfile.ZipFile(zip1, 'w') as z:
  21     z.writestr('file1', 'AAAAA')
  22     z.writestr('file2', 'BBBBB')
  23   with zipfile.ZipFile(zip2, 'w') as z:
  24     z.writestr('file2', 'ABABA' if create_conflct else 'BBBBB')
  25     z.writestr('file3', 'CCCCC')
  26   return zip1, zip2
  27
  28
  29 class ZipHelpersTest(unittest.TestCase):
  30   def test_merge_zips__identical_file(self):
  31     with tempfile.TemporaryDirectory() as tmp_dir:
  32       zip1, zip2 = _make_test_zips(tmp_dir)
  33
  34       merged_zip = os.path.join(tmp_dir, 'merged.zip')
  35       zip_helpers.merge_zips(merged_zip, [zip1, zip2])
  36
  37       with zipfile.ZipFile(merged_zip) as z:
  38         self.assertEqual(z.namelist(), ['file1', 'file2', 'file3'])
  39
  40   def test_merge_zips__conflict(self):
  41     with tempfile.TemporaryDirectory() as tmp_dir:
  42       zip1, zip2 = _make_test_zips(tmp_dir, create_conflct=True)
  43
  44       merged_zip = os.path.join(tmp_dir, 'merged.zip')
  45       with self.assertRaises(Exception):
  46         zip_helpers.merge_zips(merged_zip, [zip1, zip2])
  47
  48   def test_merge_zips__conflict_with_append(self):
  49     with tempfile.TemporaryDirectory() as tmp_dir:
  50       zip1, zip2 = _make_test_zips(tmp_dir, create_conflct=True)
  51
  52       with self.assertRaises(Exception):
  53         with zipfile.ZipFile(zip1, 'a') as dst_zip:
  54           zip_helpers.merge_zips(dst_zip, [zip2])
  55
  56
  57 if __name__ == '__main__':
  58   unittest.main()