client/bin/result_tools/dedupe_file_throttler_unittest.py

# Copyright 2017 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import os
import shutil
import tempfile
import time
import unittest

import common
from autotest_lib.client.bin.result_tools import dedupe_file_throttler
from autotest_lib.client.bin.result_tools import result_info
from autotest_lib.client.bin.result_tools import unittest_lib


# Set to 0 to force maximum throttling.
MAX_RESULT_SIZE_KB = 0

class DedupeFileThrottleTest(unittest.TestCase):
    """Test class for dedupe_file_throttler.throttle method."""

    def setUp(self):
        """Setup directory for test."""
        self.test_dir = tempfile.mkdtemp()
        self.files_to_keep = []
        self.files_to_delete = []
        modification_time = int(time.time()) - 1000

        # Files to be deduped in the root directory of result dir.
        for i in range(6):
            file_name = 'file_%d' % i
            f = os.path.join(self.test_dir, file_name)
            unittest_lib.create_file(f, unittest_lib.SIZE)
            os.utime(f, (modification_time, modification_time))
            modification_time += 1
            if (i < dedupe_file_throttler.OLDEST_FILES_TO_KEEP_COUNT or
                i >= 6 - dedupe_file_throttler.NEWEST_FILES_TO_KEEP_COUNT):
                self.files_to_keep.append(f)
            else:
                self.files_to_delete.append(f)

        folder1 = os.path.join(self.test_dir, 'folder1')
        os.mkdir(folder1)

        # Files should not be deduped.
        for i in range(3):
            file_name = 'file_not_dedupe_%d' % i
            f = os.path.join(folder1, file_name)
            unittest_lib.create_file(f, unittest_lib.SIZE)
            self.files_to_keep.append(f)

        # Files to be deduped in the sub directory of result dir.
        for i in range(10):
            file_name = 'file_in_sub_dir%d' % i
            f = os.path.join(folder1, file_name)
            unittest_lib.create_file(f, unittest_lib.SIZE)
            os.utime(f, (modification_time, modification_time))
            modification_time += 1
            if (i < dedupe_file_throttler.OLDEST_FILES_TO_KEEP_COUNT or
                i >= 10 - dedupe_file_throttler.NEWEST_FILES_TO_KEEP_COUNT):
                self.files_to_keep.append(f)
            else:
                self.files_to_delete.append(f)

    def tearDown(self):
        """Cleanup the test directory."""
        shutil.rmtree(self.test_dir, ignore_errors=True)

    def testTrim(self):
        """Test throttle method."""
        summary = result_info.ResultInfo.build_from_path(self.test_dir)
        dedupe_file_throttler.throttle(
                summary, max_result_size_KB=MAX_RESULT_SIZE_KB)

        # Verify summary sizes are updated.
        self.assertEqual(19 * unittest_lib.SIZE, summary.original_size)
        self.assertEqual(9 * unittest_lib.SIZE, summary.trimmed_size)

        # Verify files that should not be deleted still exists.
        for f in self.files_to_keep:
            self.assertTrue(os.stat(f).st_size > 0,
                            'File %s should not be deleted!' % f)

        # Verify files that should be deleted no longer exists.
        for f in self.files_to_delete:
            self.assertFalse(os.path.exists(f), 'File %s is not deleted!' % f)


# this is so the test can be run in standalone mode
if __name__ == '__main__':
    """Main"""
    unittest.main()