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 unittest
7
8import common
9from autotest_lib.client.bin.result_tools import result_info
10from autotest_lib.client.bin.result_tools import throttler_lib
11from autotest_lib.client.bin.result_tools import unittest_lib
12from autotest_lib.client.bin.result_tools import utils_lib
13
14
15# This unittest doesn't care about the size, so the size info can be shared to
16# make the code cleaner.
17FILE_SIZE_DICT = {utils_lib.ORIGINAL_SIZE_BYTES: unittest_lib.SIZE}
18
19SAMPLE_SUMMARY = {
20  '': {utils_lib.ORIGINAL_SIZE_BYTES: 7 * unittest_lib.SIZE,
21       utils_lib.DIRS: [
22         {'file1': FILE_SIZE_DICT},
23         {'file2.tar': {utils_lib.ORIGINAL_SIZE_BYTES: 2 * unittest_lib.SIZE}},
24         {'file.deleted': {utils_lib.ORIGINAL_SIZE_BYTES: unittest_lib.SIZE,
25                           utils_lib.TRIMMED_SIZE_BYTES: 0}},
26         {'keyval': FILE_SIZE_DICT},
27         {'sysinfo':
28            {utils_lib.ORIGINAL_SIZE_BYTES: 2 * unittest_lib.SIZE,
29             utils_lib.DIRS: [
30                {'file3': FILE_SIZE_DICT},
31                {'var':
32                  {utils_lib.ORIGINAL_SIZE_BYTES: unittest_lib.SIZE,
33                   utils_lib.DIRS: [
34                     {'log': {utils_lib.ORIGINAL_SIZE_BYTES: unittest_lib.SIZE,
35                              utils_lib.DIRS: [
36                                 {'file4': FILE_SIZE_DICT}
37                                 ],
38                              }
39                      }
40                    ],
41                   },
42                 }
43                ],
44             }
45            },
46          ],
47       }
48  }
49
50EXPECTED_FILES = [
51        ['', 'sysinfo', 'var', 'log', 'file4'],
52        ['', 'sysinfo', 'file3'],
53        ['', 'file2.tar'],
54        ['', 'file1'],
55        ['', 'keyval'],
56        ['', 'file.deleted'],
57        ]
58
59EXPECTED_THROTTABLE_FILES = [
60        ['', 'sysinfo', 'var', 'log', 'file4'],
61        ['', 'sysinfo', 'file3'],
62        ['', 'file2.tar'],
63        ]
64
65class ThrottlerLibTest(unittest.TestCase):
66    """Test class for methods in throttler_lib."""
67
68    def testSortResultFiles(self):
69        """Test method sort_result_files"""
70        summary = result_info.ResultInfo(parent_dir='',
71                                         original_info=SAMPLE_SUMMARY)
72        sorted_files, _ = throttler_lib.sort_result_files(summary)
73        self.assertEqual(len(EXPECTED_FILES), len(sorted_files))
74        for i in range(len(EXPECTED_FILES)):
75            self.assertEqual(os.path.join(*EXPECTED_FILES[i]),
76                             sorted_files[i].path)
77
78    def testGetThrottableFiles(self):
79        """Test method get_throttleable_files"""
80        summary = result_info.ResultInfo(parent_dir='',
81                                         original_info=SAMPLE_SUMMARY)
82        sorted_files, _ = throttler_lib.sort_result_files(summary)
83        throttleables = throttler_lib.get_throttleable_files(
84                sorted_files, ['.*file1'])
85
86        throttleables = list(throttleables)
87        self.assertEqual(len(EXPECTED_THROTTABLE_FILES), len(throttleables))
88        for i in range(len(EXPECTED_THROTTABLE_FILES)):
89            self.assertEqual(os.path.join(*EXPECTED_THROTTABLE_FILES[i]),
90                             throttleables[i].path)
91
92
93# this is so the test can be run in standalone mode
94if __name__ == '__main__':
95    """Main"""
96    unittest.main()
97