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

import unittest

import os
import split_variations_cmd

_ENABLE_FEATURES_SWITCH_NAME = 'enable-features'
_DISABLE_FEATURES_SWITCH_NAME = 'disable-features'
_FORCE_FIELD_TRIALS_SWITCH_NAME = 'force-fieldtrials'
_FORCE_FIELD_TRIAL_PARAMS_SWITCH_NAME = 'force-fieldtrial-params'


class SplitVariationsCmdUnittest(unittest.TestCase):
  def _CompareCommandLineSwitches(self, filename, cmd_list):
    """Compares two sets of command line switches.

    Args:
        filename: Name to a file that contains a set of commandline switches.
        cmd_list: A list of strings in the form of '--switch_name=switch_value'.

    Return True if they contain the same switches and each switch's values
    are the same.
    """
    assert os.path.isfile(filename)
    data = None
    with open(filename, 'r') as f:
      data = f.read().replace('\n', ' ')
    switches = split_variations_cmd.ParseCommandLineSwitchesString(data)
    if len(switches) != len(cmd_list):
      return False
    for switch_name, switch_value in switches.items():
      switch_string = '--%s="%s"' % (switch_name, switch_value)
      if switch_string not in cmd_list:
        return False
    return True

  def _GetUnittestDataDir(self):
    return os.path.join(os.path.dirname(__file__), 'unittest_data')

  def _VerifySplits(self, switch_name, splits, ref_switch_data):
    """Verifies splitting commandline switches works correctly.

    Compare that when we combine switch data from all |splits| into one,
    it's exactly the same as the |ref_switch_data|. Also check the splits are
    almost evenly distributed, that is, their data size are almost the same.

    Args:
        switch_name: The name of the switch that is verified.
        splits: A list of {switch_name: [items]} dictionaries.
                Each list element represents one of the split switch sets.
                |items| is a list of items representing switch value.
        ref_switch_data: A {switch_name: [items]} dictionary.
                         This is the switch set before splitting.
                         |items| is a list of items representing switch value.
    """
    data_lists = [
      split[switch_name] for split in splits if switch_name in split
    ]
    if len(data_lists) == 0:
      self.assertFalse(ref_switch_data)
      return
    max_size = max(len(data) for data in data_lists)
    min_size = min(len(data) for data in data_lists)
    if switch_name == _FORCE_FIELD_TRIALS_SWITCH_NAME:
      self.assertTrue(max_size - min_size <= 1)
    joined_switch_data = []
    for data in data_lists:
      joined_switch_data.extend(data)
    self.assertCountEqual(ref_switch_data, joined_switch_data)

  def testLoadFromFileAndSaveToStrings(self):
    # Verifies we load data from the file and save it to a list of strings,
    # the two data sets contain the same command line switches.
    data_file = os.path.join(self._GetUnittestDataDir(), 'variations_cmd.txt')
    assert os.path.isfile(data_file)
    data = split_variations_cmd.ParseVariationsCmdFromFile(data_file)
    cmd_list = split_variations_cmd.VariationsCmdToStrings(data)
    self.assertTrue(self._CompareCommandLineSwitches(data_file, cmd_list))

  def _testSplitVariationsCmdHelper(self, input_data):
    # Verifies we correctly and (almost) evenly split one set of command line
    # switches into two sets.
    splits = split_variations_cmd.SplitVariationsCmd(input_data)
    switches = [
      _ENABLE_FEATURES_SWITCH_NAME,
      _DISABLE_FEATURES_SWITCH_NAME,
      _FORCE_FIELD_TRIALS_SWITCH_NAME,
      _FORCE_FIELD_TRIAL_PARAMS_SWITCH_NAME,
    ]
    for switch in switches:
      self._VerifySplits(switch, splits, input_data.get(switch, []))
    # Verify both split variations are valid.
    for variations_cmd in splits:
      cmd_list = split_variations_cmd.VariationsCmdToStrings(variations_cmd)
      split_variations_cmd.ParseVariationsCmdFromString(' '.join(cmd_list))

  def testSplitVariationsCmd(self):
    input_file = os.path.join(self._GetUnittestDataDir(), 'variations_cmd.txt')
    assert os.path.isfile(input_file)
    data = split_variations_cmd.ParseVariationsCmdFromFile(input_file)
    self._testSplitVariationsCmdHelper(data)

  def testSplitVariationsCmdWithMissingEnableDisableFeatures(self):
    input_string = (
      '--force-fieldtrials="Tria1/Disabled/*Trial2/Enabled/" '
      '--force-fieldtrial-params="Trial2.Enabled:age/18/gender/male" '
      '--disable-features="FeatureA<FeatureA"'
    )
    data = split_variations_cmd.ParseVariationsCmdFromString(input_string)
    self._testSplitVariationsCmdHelper(data)

  def testSplitVariationsCmdWithMissingForceFieldTrialParams(self):
    input_string = (
      '--force-fieldtrials="*Trial2/Enabled/" '
      '--enable-features="FeatureA<FeatureA,FeatureB<FeatureB" '
      '--disable-features="FeatureC<FeatureC,FeatureD<FeatureD"'
    )
    data = split_variations_cmd.ParseVariationsCmdFromString(input_string)
    self._testSplitVariationsCmdHelper(data)

  def testSplitVariationsCmdNoFurtherSplit(self):
    input_string = (
      '--force-fieldtrials="*Trial2/Enabled/" '
      '--enable-features="FeatureA<FeatureA" '
      '--disable-features="FeatureC<FeatureC" '
      '--disable-field-trial-config'
    )
    splits = split_variations_cmd.SplitVariationsCmdFromString(input_string)
    self.assertEqual(1, len(splits))

  def testFieldTrialsEncodingParams(self):
    """Check that spaces in force-fieldtrial-params are duly decoded."""
    input_string = (
      '--force-fieldtrials="*ScaleTile MemoryLimit/Scale 120%/" '
      '--force-fieldtrial-params="ScaleTile+MemoryLimit.Scale+120%:x/100/y/Test"'
    )
    splits = split_variations_cmd.SplitVariationsCmdFromString(input_string)
    self.assertEqual(1, len(splits))

  def testSplitFeaturesGroupedByValue(self):
    """Check that features are grouped by values."""
    input_string = '--enable-features="FeatureA<Trial1,FeatureB<Trial1"'
    splits = split_variations_cmd.SplitVariationsCmdFromString(input_string)
    self.assertEqual(len(splits), 1)

  def testSplitVariationsCmdFeatureByTrials(self):
    """Check that features are assigned to corresponding trial groups"""
    input_string = (
      '--force-fieldtrials="Trial1/Disabled/*Trial2/Enabled/" '
      '--enable-features="FeatureA<Trial1,FeatureC<Trial1"'
    )
    splits = split_variations_cmd.SplitVariationsCmdFromString(input_string)
    self.assertEqual(len(splits), 2)
    for split_str in splits:
      data = split_variations_cmd.ParseVariationsCmdFromString(split_str)
      trials = [
        t.trial_name for t in data.get(_FORCE_FIELD_TRIALS_SWITCH_NAME, [])
      ]
      features = [f.key for f in data.get(_ENABLE_FEATURES_SWITCH_NAME, [])]
      if 'Trial1' in trials:
        self.assertIn('FeatureA', features)
        self.assertIn('FeatureC', features)


if __name__ == '__main__':
  unittest.main()