#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Copyright 2020 The ChromiumOS 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()