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 time
9import unittest
10
11import common
12from autotest_lib.client.bin.result_tools import dedupe_file_throttler
13from autotest_lib.client.bin.result_tools import result_info
14from autotest_lib.client.bin.result_tools import unittest_lib
15
16
17# Set to 0 to force maximum throttling.
18MAX_RESULT_SIZE_KB = 0
19
20class DedupeFileThrottleTest(unittest.TestCase):
21    """Test class for dedupe_file_throttler.throttle method."""
22
23    def setUp(self):
24        """Setup directory for test."""
25        self.test_dir = tempfile.mkdtemp()
26        self.files_to_keep = []
27        self.files_to_delete = []
28        modification_time = int(time.time()) - 1000
29
30        # Files to be deduped in the root directory of result dir.
31        for i in range(6):
32            file_name = 'file_%d' % i
33            f = os.path.join(self.test_dir, file_name)
34            unittest_lib.create_file(f, unittest_lib.SIZE)
35            os.utime(f, (modification_time, modification_time))
36            modification_time += 1
37            if (i < dedupe_file_throttler.OLDEST_FILES_TO_KEEP_COUNT or
38                i >= 6 - dedupe_file_throttler.NEWEST_FILES_TO_KEEP_COUNT):
39                self.files_to_keep.append(f)
40            else:
41                self.files_to_delete.append(f)
42
43        folder1 = os.path.join(self.test_dir, 'folder1')
44        os.mkdir(folder1)
45
46        # Files should not be deduped.
47        for i in range(3):
48            file_name = 'file_not_dedupe_%d' % i
49            f = os.path.join(folder1, file_name)
50            unittest_lib.create_file(f, unittest_lib.SIZE)
51            self.files_to_keep.append(f)
52
53        # Files to be deduped in the sub directory of result dir.
54        for i in range(10):
55            file_name = 'file_in_sub_dir%d' % i
56            f = os.path.join(folder1, file_name)
57            unittest_lib.create_file(f, unittest_lib.SIZE)
58            os.utime(f, (modification_time, modification_time))
59            modification_time += 1
60            if (i < dedupe_file_throttler.OLDEST_FILES_TO_KEEP_COUNT or
61                i >= 10 - dedupe_file_throttler.NEWEST_FILES_TO_KEEP_COUNT):
62                self.files_to_keep.append(f)
63            else:
64                self.files_to_delete.append(f)
65
66    def tearDown(self):
67        """Cleanup the test directory."""
68        shutil.rmtree(self.test_dir, ignore_errors=True)
69
70    def testTrim(self):
71        """Test throttle method."""
72        summary = result_info.ResultInfo.build_from_path(self.test_dir)
73        dedupe_file_throttler.throttle(
74                summary, max_result_size_KB=MAX_RESULT_SIZE_KB)
75
76        # Verify summary sizes are updated.
77        self.assertEqual(19 * unittest_lib.SIZE, summary.original_size)
78        self.assertEqual(9 * unittest_lib.SIZE, summary.trimmed_size)
79
80        # Verify files that should not be deleted still exists.
81        for f in self.files_to_keep:
82            self.assertTrue(os.stat(f).st_size > 0,
83                            'File %s should not be deleted!' % f)
84
85        # Verify files that should be deleted no longer exists.
86        for f in self.files_to_delete:
87            self.assertFalse(os.path.exists(f), 'File %s is not deleted!' % f)
88
89
90# this is so the test can be run in standalone mode
91if __name__ == '__main__':
92    """Main"""
93    unittest.main()
94