• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/env python3
2# -*- coding: utf-8 -*-
3# Copyright 2020 The Chromium OS Authors. All rights reserved.
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
9from __future__ import print_function
10
11import io
12from unittest.mock import patch
13import unittest
14
15from afdo_redaction import remove_cold_functions
16
17
18def _construct_profile(indices=None):
19  real_world_profile_functions = [
20      """SomeFunction1:24150:300
21 2: 75
22 3: 23850
23 39: 225
24      """,
25      """SomeFunction2:8925:225
26 0: 225
27 0.2: 150
28 0.1: SomeFunction2:6300
29  3: 6300
30 0.2: SomeFunction2:150
31  3: 75
32      """,
33      """SomeFunction3:7500:75
34 0: 75
35 0.2: 75
36 0.1: SomeFunction3:6600
37  1: 6600
38 0.2: SomeFunction3:75
39  1: 75
40      """,
41      """LargerFunction4:51450:0
42 1: 0
43 3: 0
44 3.1: 7350
45 4: 7350
46 7: 7350
47 8: 7350
48 9: 7350
49 12: 0
50 15: 0
51 13: AnotherFunction5:0
52  3: 0
53  3.1: 0
54  3.2: 0
55  4: 0
56  5: 0
57  6: 0
58  7: 0
59  8: 0
60  9: 0
61      """,
62      """SomeFakeFunction5:7500:75
63 0: 75
64 0.2: 75
65 0.1: SomeFakeFunction5:6600
66  1: 6600
67 0.2: SomeFakeFunction5:75
68  1: 75
69      """,
70  ]
71
72  ret = []
73  if not indices:
74    for x in real_world_profile_functions:
75      ret += x.strip().splitlines()
76    return ret
77
78  ret = []
79  for i in indices:
80    ret += real_world_profile_functions[i].strip().splitlines()
81  return ret
82
83
84def _run_test(input_lines, goal, cwp_file=None, benchmark_file=None):
85  input_buf = io.StringIO('\n'.join(input_lines))
86  output_buf = io.StringIO()
87  remove_cold_functions.run(input_buf, output_buf, goal, cwp_file,
88                            benchmark_file)
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  def test_remove_cold_functions_work(self):
107    input_profile_lines = _construct_profile()
108    # To make sure the cold functions are removed in order
109    expected_profile_lines = {
110        5: input_profile_lines,
111        # Entry 4 wins the tie breaker because the name is smaller
112        # alphabetically.
113        4: _construct_profile([0, 1, 3, 4]),
114        3: _construct_profile([0, 1, 3]),
115        2: _construct_profile([0, 3]),
116        1: _construct_profile([3]),
117    }
118
119    for num in expected_profile_lines:
120      self.assertCountEqual(
121          _run_test(input_profile_lines, num), expected_profile_lines[num])
122
123  def test_analyze_cwp_and_benchmark_work(self):
124    input_profile_lines = _construct_profile()
125    cwp_profile = _construct_profile([0, 1, 3, 4])
126    benchmark_profile = _construct_profile([1, 2, 3, 4])
127    cwp_buf = io.StringIO('\n'.join(cwp_profile))
128    benchmark_buf = io.StringIO('\n'.join(benchmark_profile))
129    with patch('sys.stderr', new=io.StringIO()) as fake_output:
130      _run_test(input_profile_lines, 3, cwp_buf, benchmark_buf)
131
132    output = fake_output.getvalue()
133    self.assertIn('Retained 3/5 (60.0%) functions in the profile', output)
134    self.assertIn(
135        'Retained 1/1 (100.0%) functions only appear in the CWP profile',
136        output)
137    self.assertIn(
138        'Retained 0/1 (0.0%) functions only appear in the benchmark profile',
139        output)
140    self.assertIn(
141        'Retained 2/3 (66.7%) functions appear in both CWP and benchmark'
142        ' profiles', output)
143
144
145if __name__ == '__main__':
146  unittest.main()
147