1 files changed, 146 insertions, 0 deletions
diff --git a/afdo_redaction/remove_cold_functions_test.py b/afdo_redaction/remove_cold_functions_test.py
new file mode 100755
index 00000000..14f946b0
--- /dev/null
+++ b/afdo_redaction/remove_cold_functions_test.py
@@ -0,0 +1,146 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Copyright 2020 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.
+
+"""Tests for remove_cold_functions."""
+
+from __future__ import print_function
+
+import io
+from unittest.mock import patch
+import unittest
+
+from afdo_redaction import remove_cold_functions
+
+
+def _construct_profile(indices=None):
+  real_world_profile_functions = [
+      """SomeFunction1:24150:300
+ 2: 75
+ 3: 23850
+ 39: 225
+      """,
+      """SomeFunction2:8925:225
+ 0: 225
+ 0.2: 150
+ 0.1: SomeFunction2:6300
+  3: 6300
+ 0.2: SomeFunction2:150
+  3: 75
+      """,
+      """SomeFunction3:7500:75
+ 0: 75
+ 0.2: 75
+ 0.1: SomeFunction3:6600
+  1: 6600
+ 0.2: SomeFunction3:75
+  1: 75
+      """,
+      """LargerFunction4:51450:0
+ 1: 0
+ 3: 0
+ 3.1: 7350
+ 4: 7350
+ 7: 7350
+ 8: 7350
+ 9: 7350
+ 12: 0
+ 15: 0
+ 13: AnotherFunction5:0
+  3: 0
+  3.1: 0
+  3.2: 0
+  4: 0
+  5: 0
+  6: 0
+  7: 0
+  8: 0
+  9: 0
+      """,
+      """SomeFakeFunction5:7500:75
+ 0: 75
+ 0.2: 75
+ 0.1: SomeFakeFunction5:6600
+  1: 6600
+ 0.2: SomeFakeFunction5:75
+  1: 75
+      """,
+  ]
+
+  ret = []
+  if not indices:
+    for x in real_world_profile_functions:
+      ret += x.strip().splitlines()
+    return ret
+
+  ret = []
+  for i in indices:
+    ret += real_world_profile_functions[i].strip().splitlines()
+  return ret
+
+
+def _run_test(input_lines, goal, cwp_file=None, benchmark_file=None):
+  input_buf = io.StringIO('\n'.join(input_lines))
+  output_buf = io.StringIO()
+  remove_cold_functions.run(input_buf, output_buf, goal, cwp_file,
+                            benchmark_file)
+  return output_buf.getvalue().splitlines()
+
+
+class Test(unittest.TestCase):
+  """Test functions in remove_cold_functions.py"""
+
+  def test_empty_profile(self):
+    self.assertEqual(_run_test([], 0), [])
+
+  def test_remove_all_functions_fail(self):
+    input_profile_lines = _construct_profile()
+    with self.assertRaises(Exception) as context:
+      _run_test(input_profile_lines, 0)
+    self.assertEqual(
+        str(context.exception),
+        "It's invalid to remove all functions in the profile")
+
+  def test_remove_cold_functions_work(self):
+    input_profile_lines = _construct_profile()
+    # To make sure the cold functions are removed in order
+    expected_profile_lines = {
+        5: input_profile_lines,
+        # Entry 4 wins the tie breaker because the name is smaller
+        # alphabetically.
+        4: _construct_profile([0, 1, 3, 4]),
+        3: _construct_profile([0, 1, 3]),
+        2: _construct_profile([0, 3]),
+        1: _construct_profile([3]),
+    }
+
+    for num in expected_profile_lines:
+      self.assertCountEqual(
+          _run_test(input_profile_lines, num), expected_profile_lines[num])
+
+  def test_analyze_cwp_and_benchmark_work(self):
+    input_profile_lines = _construct_profile()
+    cwp_profile = _construct_profile([0, 1, 3, 4])
+    benchmark_profile = _construct_profile([1, 2, 3, 4])
+    cwp_buf = io.StringIO('\n'.join(cwp_profile))
+    benchmark_buf = io.StringIO('\n'.join(benchmark_profile))
+    with patch('sys.stderr', new=io.StringIO()) as fake_output:
+      _run_test(input_profile_lines, 3, cwp_buf, benchmark_buf)
+
+    output = fake_output.getvalue()
+    self.assertIn('Retained 3/5 (60.0%) functions in the profile', output)
+    self.assertIn(
+        'Retained 1/1 (100.0%) functions only appear in the CWP profile',
+        output)
+    self.assertIn(
+        'Retained 0/1 (0.0%) functions only appear in the benchmark profile',
+        output)
+    self.assertIn(
+        'Retained 2/3 (66.7%) functions appear in both CWP and benchmark'
+        ' profiles', output)
+
+
+if __name__ == '__main__':
+  unittest.main()