"""Tests for devicemap_lib.py module."""

import unittest
import yaml

from catatester.devicemap import devicemap_lib

_SAMPLE_YAML = """
allowed_dimensions:
  Device:
    - ip
    - device_ip
hosts:
  host-1:
    host_ip: 1.2.3.4
    location: 'US-MTV-B47'
    hosted_devices:
      - device: {ip: 192.168.0.1}
      - device: {ip: 192.168.0.2}
      - device:
          ip: 192.168.0.3
          location: 'NOT US-MTV-B47'
"""


class DeviceMapLibTest(unittest.TestCase):
  """Tests for devicemap_lib.py module."""

  def setUp(self):
    self.sample_configs = devicemap_lib.Configs(yaml.safe_load(_SAMPLE_YAML))

  def assertBotConfigMatchesSwarmingSchema(self, bot_config):
    config_dict = bot_config.as_dict
    self.assertTrue('dimensions' in config_dict)
    self.assertTrue('state' in config_dict)
    for value in config_dict['dimensions'].values():
      self.assertTrue(hasattr(value, '__iter__'))

  def testReadDeviceMap(self):
    self.assertTrue(
      devicemap_lib._ReadDeviceMap())

  def testGetNumBotConfigs(self):
    self.assertEqual(3, self.sample_configs.GetNumBotConfigs('host-1'))
    with self.assertRaises(devicemap_lib.DeviceMapError):
      self.sample_configs.GetNumBotConfigs('undefined-host')

  def testGetAllowedDimensions(self):
    expected = set(['ip', 'device_ip'])
    actual = self.sample_configs.GetAllowedDimensions('Device')
    self.assertEqual(expected, actual)
    with self.assertRaises(devicemap_lib.DeviceMapError):
      self.sample_configs.GetAllowedDimensions('undefined-type')

  def testGetBotConfig_canGetDevice(self):
    bot_config = self.sample_configs.GetBotConfig('host-1', 1)
    self.assertTrue(bot_config)
    self.assertBotConfigMatchesSwarmingSchema(bot_config)
    self.assertIn('location', bot_config.dimensions)
    self.assertEqual('US-MTV-B47', bot_config.dimensions['location'])
    # Also verify works with a slightly different format (newline instead of {})
    bot_config = self.sample_configs.GetBotConfig('host-1', 3)
    self.assertTrue(bot_config)
    self.assertBotConfigMatchesSwarmingSchema(bot_config)
    self.assertTrue('device_type' in bot_config.dimensions)
    self.assertEqual('device', bot_config.dimensions['device_type'])
    self.assertTrue('device_ip' in bot_config.dimensions)
    self.assertEqual('192.168.0.3', bot_config.dimensions['device_ip'])
    self.assertIn('location', bot_config.dimensions)
    self.assertEqual('NOT US-MTV-B47', bot_config.dimensions['location'])

  def testGetBotConfig_canGetDeviceWithNewDimension(self):
    sample_yaml = """
hosts:
  host-2:
    host_ip: 2.2.2.2
    hosted_devices:
      - device: {ip: 192.168.0.4, new_dimension: hello}"""
    configs = devicemap_lib.Configs(yaml.safe_load(sample_yaml))
    bot_config = configs.GetBotConfig('host-2', 1)
    self.assertTrue(bot_config)
    self.assertBotConfigMatchesSwarmingSchema(bot_config)
    expected = ('new_dimension', 'hello')
    self.assertTrue(expected in bot_config.dimensions.items())

  def testGetBotConfig_canGetDeviceWithDimensionContainingAList(self):
    sample_yaml = """
hosts:
  host-2:
    host_ip: 2.2.2.2
    hosted_devices:
      - device: {ip: 192.168.0.4, new_dimension: ['hello']}"""
    configs = devicemap_lib.Configs(yaml.safe_load(sample_yaml))
    bot_config = configs.GetBotConfig('host-2', 1)
    self.assertTrue(bot_config)
    self.assertBotConfigMatchesSwarmingSchema(bot_config)
    expected = ('new_dimension', ['hello'])
    self.assertTrue(expected in bot_config.dimensions.items())

  def testGetBotConfig_canGetDeviceWithExtraState(self):
    sample_yaml = """
hosts:
  host-2:
    host_ip: 2.2.2.2
    hosted_devices:
      - device: {ip: 192.168.0.5, extra_state: {eth_switch: 100.107.119.209}}
      - device:
          ip: 192.168.0.6
          extra_state:
            eth_switch: 100.107.119.210
            example: 'something new'"""
    configs = devicemap_lib.Configs(yaml.safe_load(sample_yaml))
    bot_config = configs.GetBotConfig('host-2', 1)
    self.assertTrue(bot_config)
    self.assertBotConfigMatchesSwarmingSchema(bot_config)
    expected = ('eth_switch', '100.107.119.209')
    self.assertTrue(expected in bot_config.state.items())

    # Also verify works with a slightly different format (newline instead of {})
    bot_config = configs.GetBotConfig('host-2', 2)
    self.assertTrue(bot_config)
    self.assertBotConfigMatchesSwarmingSchema(bot_config)
    expected = ('eth_switch', '100.107.119.210')
    self.assertTrue(expected in bot_config.state.items())

  def testGetBotConfig_canGetTestbed(self):
    sample_yaml = """
hosts:
  host-2:
    host_ip: 2.2.2.2
    location: 'US-MTV-B45'
    hosted_devices:
      - Testbed:
          dimensions:
            device_type: chaos_testbed
            device_ip: 172.22.38.73
            group: ['chaos_lab']
          extra_state:
            anchovy: ['172.22.38.69', '172.22.38.70']
            biggie: '172.22.38.73'
            chorizo: ['172.22.38.66', '172.22.38.67']
            ethernet_switch: '172.22.38.80'
            linux_client: '172.22.38.81'
            mushroom: '172.22.38.72'
            pepperoni: '172.22.38.71'
            pineapple: ['172.22.38.78', '172.22.38.68']
            rpi: '172.22.38.79'
            salami: ['172.22.38.65', '172.22.38.83']
            steak: ['172.22.38.74', '172.22.38.64']"""
    configs = devicemap_lib.Configs(yaml.safe_load(sample_yaml))
    bot_config = configs.GetBotConfig('host-2', 1)
    self.assertTrue(bot_config)
    self.assertBotConfigMatchesSwarmingSchema(bot_config)
    expected = ('biggie', '172.22.38.73')
    self.assertTrue(expected in bot_config.state.items())
    expected = ('pineapple', ['172.22.38.78', '172.22.38.68'])
    self.assertTrue(expected in bot_config.state.items())
    self.assertIn('location', bot_config.dimensions)
    self.assertEqual('US-MTV-B45', bot_config.dimensions['location'])

  def testGetBotConfig_undefinedHostRaisesError(self):
    with self.assertRaises(devicemap_lib.DeviceMapError):
      self.sample_configs.GetBotConfig('host-doesnt-exist', 1)

  def testGetBotConfig_undefinedDeviceRaisesError(self):
    with self.assertRaises(devicemap_lib.DeviceMapError):
      self.sample_configs.GetBotConfig('host-1', 9999)

  def testFactoryResetNoSignedInGroup(self):
    """Test that FDR flag is set when device is not in the group 'signed_in'."""
    sample_yaml = """
hosts:
  host-2:
    host_ip: 2.2.2.2
    hosted_devices:
      - device:
          ip: 100.127.32.17
          group: ['device_performance']"""
    configs = devicemap_lib.Configs(yaml.safe_load(sample_yaml))
    bot_config = configs.GetBotConfig('host-2', 1)
    self.assertTrue(bot_config)
    self.assertEqual('true', bot_config.dimensions['factory_reset'])

  def testFactoryResetNoGroup(self):
    """Test that FDR flag is set when device is not explicitly in a group."""
    sample_yaml = """
hosts:
  host-2:
    host_ip: 2.2.2.2
    hosted_devices:
      - device:
          ip: 100.127.32.17"""
    configs = devicemap_lib.Configs(yaml.safe_load(sample_yaml))
    bot_config = configs.GetBotConfig('host-2', 1)
    self.assertTrue(bot_config)
    self.assertEqual('true', bot_config.dimensions['factory_reset'])

  def testFactoryResetSignedInGroup(self):
    """Test that FDR flag is set when device is in the group 'signed_in'."""
    sample_yaml = """
hosts:
  host-2:
    host_ip: 2.2.2.2
    hosted_devices:
      - device:
          ip: 100.127.32.17
          group: ['signed_in', 'device_performance']"""
    configs = devicemap_lib.Configs(yaml.safe_load(sample_yaml))
    bot_config = configs.GetBotConfig('host-2', 1)
    self.assertTrue(bot_config)
    self.assertEqual('false', bot_config.dimensions['factory_reset'])

  def testFactoryResetExplicit(self):
    """Test that FDR flag is set explicitly specified in config."""
    sample_yaml = """
hosts:
  host-2:
    host_ip: 2.2.2.2
    hosted_devices:
      - device:
          ip: 100.127.32.17
          factory_reset: 'false'"""
    configs = devicemap_lib.Configs(yaml.safe_load(sample_yaml))
    bot_config = configs.GetBotConfig('host-2', 1)
    self.assertTrue(bot_config)
    self.assertEqual('false', bot_config.dimensions['factory_reset'])


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