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
5import os
6import shutil
7import tempfile
8import unittest
9
10import common
11from autotest_lib.client.bin.result_tools import delete_file_throttler
12from autotest_lib.client.bin.result_tools import result_info
13from autotest_lib.client.bin.result_tools import unittest_lib
14from autotest_lib.client.bin.result_tools import utils_lib
15
16
17LARGE_SIZE_BYTE = 1000
18MEDIUM_SIZE_BYTE = 800
19SMALL_SIZE_BYTE = 100
20# Maximum result size is set to 3KB so the file with MEDIUM_SIZE_BYTE will be
21# kept.
22MAX_RESULT_SIZE_KB = 5
23# Any file with size above the threshold is qualified to be deleted.
24FILE_SIZE_THRESHOLD_BYTE = 512
25
26SUMMARY_AFTER_THROTTLE = {
27    '': {utils_lib.DIRS: [
28            {'AndroidDeviceXXX': {
29                utils_lib.DIRS: [
30                    {'file5': {utils_lib.ORIGINAL_SIZE_BYTES: LARGE_SIZE_BYTE}},
31                    ],
32                utils_lib.ORIGINAL_SIZE_BYTES: LARGE_SIZE_BYTE}},
33            {'chrome.123.perf.data': {
34                utils_lib.ORIGINAL_SIZE_BYTES: LARGE_SIZE_BYTE}},
35            {'file1.xml': {utils_lib.ORIGINAL_SIZE_BYTES: LARGE_SIZE_BYTE,
36                          utils_lib.TRIMMED_SIZE_BYTES: 0}},
37            {'file2.jpg': {utils_lib.ORIGINAL_SIZE_BYTES: LARGE_SIZE_BYTE,
38                          utils_lib.TRIMMED_SIZE_BYTES: 0}},
39            {'file3.log': {utils_lib.ORIGINAL_SIZE_BYTES: SMALL_SIZE_BYTE}},
40            {'file_to_keep': {utils_lib.ORIGINAL_SIZE_BYTES: MEDIUM_SIZE_BYTE}},
41            {'folder1': {
42                utils_lib.DIRS: [
43                    {'file4': {utils_lib.ORIGINAL_SIZE_BYTES: LARGE_SIZE_BYTE,
44                              utils_lib.TRIMMED_SIZE_BYTES: 0}},
45                    {'keyval':
46                        {utils_lib.ORIGINAL_SIZE_BYTES: LARGE_SIZE_BYTE}},
47                    ],
48                utils_lib.ORIGINAL_SIZE_BYTES: 2 * LARGE_SIZE_BYTE,
49                utils_lib.TRIMMED_SIZE_BYTES: LARGE_SIZE_BYTE}},
50            {'test_run_details.txt': {
51                utils_lib.ORIGINAL_SIZE_BYTES: LARGE_SIZE_BYTE}}],
52         utils_lib.ORIGINAL_SIZE_BYTES:
53                7 * LARGE_SIZE_BYTE + SMALL_SIZE_BYTE + MEDIUM_SIZE_BYTE,
54         utils_lib.TRIMMED_SIZE_BYTES:
55                4 * LARGE_SIZE_BYTE + SMALL_SIZE_BYTE + MEDIUM_SIZE_BYTE}
56    }
57
58class ThrottleTest(unittest.TestCase):
59    """Test class for shrink_file_throttler.throttle method."""
60
61    def setUp(self):
62        """Setup directory for test."""
63        self.test_dir = tempfile.mkdtemp()
64        self.files_not_deleted = []
65        self.files_to_delete = []
66
67        file1 = os.path.join(self.test_dir, 'file1.xml')
68        unittest_lib.create_file(file1, LARGE_SIZE_BYTE)
69        self.files_to_delete.append(file1)
70
71        file2 = os.path.join(self.test_dir, 'file2.jpg')
72        unittest_lib.create_file(file2, LARGE_SIZE_BYTE)
73        self.files_to_delete.append(file2)
74
75        file_to_keep = os.path.join(self.test_dir, 'file_to_keep')
76        unittest_lib.create_file(file_to_keep, MEDIUM_SIZE_BYTE)
77        self.files_not_deleted.append(file_to_keep)
78
79        file3 = os.path.join(self.test_dir, 'file3.log')
80        unittest_lib.create_file(file3, SMALL_SIZE_BYTE)
81        self.files_not_deleted.append(file3)
82
83        folder1 = os.path.join(self.test_dir, 'folder1')
84        os.mkdir(folder1)
85        file4 = os.path.join(folder1, 'file4')
86        unittest_lib.create_file(file4, LARGE_SIZE_BYTE)
87        self.files_to_delete.append(file4)
88
89        protected_file = os.path.join(folder1, 'keyval')
90        unittest_lib.create_file(protected_file, LARGE_SIZE_BYTE)
91        self.files_not_deleted.append(protected_file)
92
93        folder2 = os.path.join(self.test_dir, 'AndroidDeviceXXX')
94        os.mkdir(folder2)
95        file5 = os.path.join(folder2, 'file5')
96        unittest_lib.create_file(file5, LARGE_SIZE_BYTE)
97        self.files_not_deleted.append(file5)
98
99        test_run = os.path.join(self.test_dir, 'test_run_details.txt')
100        unittest_lib.create_file(test_run, LARGE_SIZE_BYTE)
101        self.files_not_deleted.append(test_run)
102
103        perf_data = os.path.join(self.test_dir, 'chrome.123.perf.data')
104        unittest_lib.create_file(perf_data, LARGE_SIZE_BYTE)
105        self.files_not_deleted.append(perf_data)
106
107    def tearDown(self):
108        """Cleanup the test directory."""
109        shutil.rmtree(self.test_dir, ignore_errors=True)
110
111    def testTrim(self):
112        """Test throttle method."""
113        summary = result_info.ResultInfo.build_from_path(self.test_dir)
114        delete_file_throttler.throttle(
115                summary,
116                max_result_size_KB=MAX_RESULT_SIZE_KB,
117                file_size_threshold_byte=FILE_SIZE_THRESHOLD_BYTE)
118
119        self.assertEqual(SUMMARY_AFTER_THROTTLE, summary)
120
121        # Verify files that should not be deleted still exists.
122        for f in self.files_not_deleted:
123            self.assertTrue(os.stat(f).st_size > 0,
124                            'File %s should not be deleted!' % f)
125
126        # Verify files that should be deleted no longer exists.
127        for f in self.files_to_delete:
128            self.assertFalse(os.path.exists(f), 'File %s is not deleted!' % f)
129
130
131# this is so the test can be run in standalone mode
132if __name__ == '__main__':
133    """Main"""
134    unittest.main()