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