1 # Copyright 2017 The Chromium OS Authors. All rights reserved. 2 # Use of this source code is governed by a BSD-style license that can be 3 # found in the LICENSE file. 4 5 import os 6 import shutil 7 import tempfile 8 import unittest 9 10 import common 11 from autotest_lib.client.bin.result_tools import result_info 12 from autotest_lib.client.bin.result_tools import shrink_file_throttler 13 from autotest_lib.client.bin.result_tools import unittest_lib 14 from autotest_lib.client.bin.result_tools import utils_lib 15 16 17 ORIGINAL_SIZE_BYTE = 1000 18 MAX_RESULT_SIZE_KB = 2 19 FILE_SIZE_LIMIT_BYTE = 512 20 LARGE_SIZE_BYTE = 100 * 1024 21 22 SUMMARY_AFTER_TRIMMING = { 23 '': {utils_lib.DIRS: [ 24 {'BUILD_INFO-HT7591A00171': { 25 utils_lib.ORIGINAL_SIZE_BYTES: ORIGINAL_SIZE_BYTE}}, 26 {'file1.xml': {utils_lib.ORIGINAL_SIZE_BYTES: ORIGINAL_SIZE_BYTE}}, 27 {'file2.jpg': {utils_lib.ORIGINAL_SIZE_BYTES: ORIGINAL_SIZE_BYTE}}, 28 {'file3.log': {utils_lib.ORIGINAL_SIZE_BYTES: ORIGINAL_SIZE_BYTE, 29 utils_lib.TRIMMED_SIZE_BYTES: FILE_SIZE_LIMIT_BYTE}}, 30 {'folder1': { 31 utils_lib.DIRS: [ 32 {'file4': { 33 utils_lib.ORIGINAL_SIZE_BYTES: ORIGINAL_SIZE_BYTE, 34 utils_lib.TRIMMED_SIZE_BYTES: FILE_SIZE_LIMIT_BYTE}}], 35 utils_lib.ORIGINAL_SIZE_BYTES: ORIGINAL_SIZE_BYTE, 36 utils_lib.TRIMMED_SIZE_BYTES: FILE_SIZE_LIMIT_BYTE}}, 37 {'test_run_details.txt': { 38 utils_lib.ORIGINAL_SIZE_BYTES: ORIGINAL_SIZE_BYTE}}], 39 utils_lib.ORIGINAL_SIZE_BYTES: 6 * ORIGINAL_SIZE_BYTE, 40 utils_lib.TRIMMED_SIZE_BYTES: ( 41 4 * ORIGINAL_SIZE_BYTE + 2 * FILE_SIZE_LIMIT_BYTE)} 42 } 43 44 OLD_TIME = 1498800000 45 46 class ShrinkFileThrottleTest(unittest.TestCase): 47 """Test class for shrink_file_throttler.throttle method.""" 48 49 def setUp(self): 50 """Setup directory for test.""" 51 self.test_dir = tempfile.mkdtemp() 52 self.files_not_shrink = [] 53 self.files_to_shrink = [] 54 55 build_info = os.path.join(self.test_dir, 'BUILD_INFO-HT7591A00171') 56 unittest_lib.create_file(build_info, ORIGINAL_SIZE_BYTE) 57 self.files_not_shrink.append(build_info) 58 59 file1 = os.path.join(self.test_dir, 'file1.xml') 60 unittest_lib.create_file(file1, ORIGINAL_SIZE_BYTE) 61 self.files_not_shrink.append(file1) 62 63 file2 = os.path.join(self.test_dir, 'file2.jpg') 64 unittest_lib.create_file(file2, ORIGINAL_SIZE_BYTE) 65 self.files_not_shrink.append(file2) 66 67 file3 = os.path.join(self.test_dir, 'file3.log') 68 unittest_lib.create_file(file3, ORIGINAL_SIZE_BYTE) 69 self.files_to_shrink.append(file3) 70 os.utime(file3, (OLD_TIME, OLD_TIME)) 71 72 file4 = os.path.join(self.test_dir, 'test_run_details.txt') 73 unittest_lib.create_file(file4, ORIGINAL_SIZE_BYTE) 74 self.files_not_shrink.append(file4) 75 76 folder1 = os.path.join(self.test_dir, 'folder1') 77 os.mkdir(folder1) 78 file4 = os.path.join(folder1, 'file4') 79 unittest_lib.create_file(file4, ORIGINAL_SIZE_BYTE) 80 self.files_to_shrink.append(file4) 81 os.utime(file4, (OLD_TIME, OLD_TIME)) 82 83 def tearDown(self): 84 """Cleanup the test directory.""" 85 shutil.rmtree(self.test_dir, ignore_errors=True) 86 87 def testTrim(self): 88 """Test throttle method.""" 89 summary = result_info.ResultInfo.build_from_path(self.test_dir) 90 shrink_file_throttler.throttle( 91 summary, 92 max_result_size_KB=MAX_RESULT_SIZE_KB, 93 file_size_limit_byte=FILE_SIZE_LIMIT_BYTE) 94 95 self.assertEqual(SUMMARY_AFTER_TRIMMING, summary) 96 97 # Verify files that should not be shrunk are not changed. 98 for f in self.files_not_shrink: 99 self.assertEqual(ORIGINAL_SIZE_BYTE, os.stat(f).st_size, 100 'File %s should not be shrank!' % f) 101 102 # Verify files that should be shrunk are updated. 103 for f in self.files_to_shrink: 104 stat = os.stat(f) 105 self.assertTrue(FILE_SIZE_LIMIT_BYTE >= stat.st_size, 106 'File %s is not shrank!' % f) 107 self.assertEqual(OLD_TIME, stat.st_mtime) 108 109 110 class MultipleShrinkFileTest(unittest.TestCase): 111 """Test class for shrink_file_throttler.throttle method for files to be 112 shrunk multiple times. 113 """ 114 115 def setUp(self): 116 """Setup directory for test.""" 117 self.test_dir = tempfile.mkdtemp() 118 119 self.file_to_shrink = os.path.join(self.test_dir, 'file1.txt') 120 unittest_lib.create_file(self.file_to_shrink, LARGE_SIZE_BYTE) 121 122 def tearDown(self): 123 """Cleanup the test directory.""" 124 shutil.rmtree(self.test_dir, ignore_errors=True) 125 126 def testTrim(self): 127 """Shrink the file twice and check its content.""" 128 summary = result_info.ResultInfo.build_from_path( 129 parent_dir=self.test_dir, name=utils_lib.ROOT_DIR, 130 top_dir=self.test_dir, parent_result_info=None) 131 shrink_file_throttler.throttle( 132 summary, max_result_size_KB=60, 133 file_size_limit_byte=50 * 1024) 134 summary = result_info.ResultInfo.build_from_path(self.test_dir) 135 shrink_file_throttler.throttle( 136 summary, max_result_size_KB=30, 137 file_size_limit_byte=25 * 1024) 138 139 with open(self.file_to_shrink) as f: 140 content = f.read() 141 142 original_size_string = (shrink_file_throttler.ORIGINAL_SIZE_TEMPLATE % 143 LARGE_SIZE_BYTE) 144 self.assertTrue(original_size_string in content) 145 146 trimmed_size_string = ( 147 shrink_file_throttler.TRIMMED_FILE_INJECT_TEMPLATE % 76990) 148 self.assertTrue(trimmed_size_string in content) 149 150 151 # this is so the test can be run in standalone mode 152 if __name__ == '__main__': 153 """Main""" 154 unittest.main() 155
