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