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