1#!/usr/bin/env python3 2# -*- coding: utf-8 -*- 3# Copyright 2020 The ChromiumOS Authors 4# Use of this source code is governed by a BSD-style license that can be 5# found in the LICENSE file. 6 7"""Tests for remove_cold_functions.""" 8 9 10import io 11import unittest 12from unittest.mock import patch 13 14from afdo_redaction import remove_cold_functions 15 16 17def _construct_profile(indices=None): 18 real_world_profile_functions = [ 19 """SomeFunction1:24150:300 20 2: 75 21 3: 23850 22 39: 225 23 """, 24 """SomeFunction2:8925:225 25 0: 225 26 0.2: 150 27 0.1: SomeFunction2:6300 28 3: 6300 29 0.2: SomeFunction2:150 30 3: 75 31 """, 32 """SomeFunction3:7500:75 33 0: 75 34 0.2: 75 35 0.1: SomeFunction3:6600 36 1: 6600 37 0.2: SomeFunction3:75 38 1: 75 39 """, 40 """LargerFunction4:51450:0 41 1: 0 42 3: 0 43 3.1: 7350 44 4: 7350 45 7: 7350 46 8: 7350 47 9: 7350 48 12: 0 49 15: 0 50 13: AnotherFunction5:0 51 3: 0 52 3.1: 0 53 3.2: 0 54 4: 0 55 5: 0 56 6: 0 57 7: 0 58 8: 0 59 9: 0 60 """, 61 """SomeFakeFunction5:7500:75 62 0: 75 63 0.2: 75 64 0.1: SomeFakeFunction5:6600 65 1: 6600 66 0.2: SomeFakeFunction5:75 67 1: 75 68 """, 69 ] 70 71 ret = [] 72 if not indices: 73 for x in real_world_profile_functions: 74 ret += x.strip().splitlines() 75 return ret 76 77 ret = [] 78 for i in indices: 79 ret += real_world_profile_functions[i].strip().splitlines() 80 return ret 81 82 83def _run_test(input_lines, goal, cwp_file=None, benchmark_file=None): 84 input_buf = io.StringIO("\n".join(input_lines)) 85 output_buf = io.StringIO() 86 remove_cold_functions.run( 87 input_buf, output_buf, goal, cwp_file, benchmark_file 88 ) 89 return output_buf.getvalue().splitlines() 90 91 92class Test(unittest.TestCase): 93 """Test functions in remove_cold_functions.py""" 94 95 def test_empty_profile(self): 96 self.assertEqual(_run_test([], 0), []) 97 98 def test_remove_all_functions_fail(self): 99 input_profile_lines = _construct_profile() 100 with self.assertRaises(Exception) as context: 101 _run_test(input_profile_lines, 0) 102 self.assertEqual( 103 str(context.exception), 104 "It's invalid to remove all functions in the profile", 105 ) 106 107 def test_remove_cold_functions_work(self): 108 input_profile_lines = _construct_profile() 109 # To make sure the cold functions are removed in order 110 expected_profile_lines = { 111 5: input_profile_lines, 112 # Entry 4 wins the tie breaker because the name is smaller 113 # alphabetically. 114 4: _construct_profile([0, 1, 3, 4]), 115 3: _construct_profile([0, 1, 3]), 116 2: _construct_profile([0, 3]), 117 1: _construct_profile([3]), 118 } 119 120 for num in expected_profile_lines: 121 self.assertCountEqual( 122 _run_test(input_profile_lines, num), expected_profile_lines[num] 123 ) 124 125 def test_analyze_cwp_and_benchmark_work(self): 126 input_profile_lines = _construct_profile() 127 cwp_profile = _construct_profile([0, 1, 3, 4]) 128 benchmark_profile = _construct_profile([1, 2, 3, 4]) 129 cwp_buf = io.StringIO("\n".join(cwp_profile)) 130 benchmark_buf = io.StringIO("\n".join(benchmark_profile)) 131 with patch("sys.stderr", new=io.StringIO()) as fake_output: 132 _run_test(input_profile_lines, 3, cwp_buf, benchmark_buf) 133 134 output = fake_output.getvalue() 135 self.assertIn("Retained 3/5 (60.0%) functions in the profile", output) 136 self.assertIn( 137 "Retained 1/1 (100.0%) functions only appear in the CWP profile", 138 output, 139 ) 140 self.assertIn( 141 "Retained 0/1 (0.0%) functions only appear in the benchmark profile", 142 output, 143 ) 144 self.assertIn( 145 "Retained 2/3 (66.7%) functions appear in both CWP and benchmark" 146 " profiles", 147 output, 148 ) 149 150 151if __name__ == "__main__": 152 unittest.main() 153