lib/cros_collections_unittest.py

# -*- coding: utf-8 -*-
# Copyright 2018 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.

"""Test the cros_collections module."""

from __future__ import print_function

import collections

from chromite.lib import cros_collections
from chromite.lib import cros_test_lib


class CollectionTest(cros_test_lib.TestCase):
  """Tests for Collection helper."""

  def testDefaults(self):
    """Verify default values kick in."""
    O = cros_collections.Collection('O', a=0, b='string', c={})
    o = O()
    self.assertEqual(o.a, 0)
    self.assertEqual(o.b, 'string')
    self.assertEqual(o.c, {})

  def testOverrideDefaults(self):
    """Verify we can set custom values at instantiation time."""
    O = cros_collections.Collection('O', a=0, b='string', c={})
    o = O(a=1000)
    self.assertEqual(o.a, 1000)
    self.assertEqual(o.b, 'string')
    self.assertEqual(o.c, {})

  def testSetNoNewMembers(self):
    """Verify we cannot add new members after the fact."""
    O = cros_collections.Collection('O', a=0, b='string', c={})
    o = O()

    # Need the func since self.assertRaises evaluates the args in this scope.
    def _setit(collection):
      collection.does_not_exit = 10
    self.assertRaises(AttributeError, _setit, o)
    self.assertRaises(AttributeError, setattr, o, 'new_guy', 10)

  def testGetNoNewMembers(self):
    """Verify we cannot get new members after the fact."""
    O = cros_collections.Collection('O', a=0, b='string', c={})
    o = O()

    # Need the func since self.assertRaises evaluates the args in this scope.
    def _getit(collection):
      return collection.does_not_exit
    self.assertRaises(AttributeError, _getit, o)
    self.assertRaises(AttributeError, getattr, o, 'foooo')

  def testNewValue(self):
    """Verify we change members correctly."""
    O = cros_collections.Collection('O', a=0, b='string', c={})
    o = O()
    o.a = 'a string'
    o.c = 123
    self.assertEqual(o.a, 'a string')
    self.assertEqual(o.b, 'string')
    self.assertEqual(o.c, 123)

  def testString(self):
    """Make sure the string representation is readable by da hue mans."""
    O = cros_collections.Collection('O', a=0, b='string', c={})
    o = O()
    self.assertEqual("Collection_O(a=0, b='string', c={})", str(o))


class TestGroupByKey(cros_test_lib.TestCase):
  """Test SplitByKey."""

  def testEmpty(self):
    self.assertEqual({}, cros_collections.GroupByKey([], ''))

  def testGroupByKey(self):
    input_iter = [{'a': None, 'b': 0},
                  {'a': 1, 'b': 1},
                  {'a': 2, 'b': 2},
                  {'a': 1, 'b': 3}]
    expected_result = {
        None: [{'a': None, 'b': 0}],
        1:    [{'a': 1, 'b': 1},
               {'a': 1, 'b': 3}],
        2:    [{'a': 2, 'b': 2}]}
    self.assertEqual(cros_collections.GroupByKey(input_iter, 'a'),
                     expected_result)


class GroupNamedtuplesByKeyTests(cros_test_lib.TestCase):
  """Tests for GroupNamedtuplesByKey"""

  def testGroupNamedtuplesByKeyWithEmptyInputIter(self):
    """Test GroupNamedtuplesByKey with empty input_iter."""
    self.assertEqual({}, cros_collections.GroupByKey([], ''))

  def testGroupNamedtuplesByKey(self):
    """Test GroupNamedtuplesByKey."""
    TestTuple = collections.namedtuple('TestTuple', ('key1', 'key2'))
    r1 = TestTuple('t1', 'val1')
    r2 = TestTuple('t2', 'val2')
    r3 = TestTuple('t2', 'val2')
    r4 = TestTuple('t3', 'val3')
    r5 = TestTuple('t3', 'val3')
    r6 = TestTuple('t3', 'val3')
    input_iter = [r1, r2, r3, r4, r5, r6]

    expected_result = {
        't1': [r1],
        't2': [r2, r3],
        't3': [r4, r5, r6]}
    self.assertDictEqual(
        cros_collections.GroupNamedtuplesByKey(input_iter, 'key1'),
        expected_result)

    expected_result = {
        'val1': [r1],
        'val2': [r2, r3],
        'val3': [r4, r5, r6]}
    self.assertDictEqual(
        cros_collections.GroupNamedtuplesByKey(input_iter, 'key2'),
        expected_result)

    expected_result = {
        None: [r1, r2, r3, r4, r5, r6]}
    self.assertDictEqual(
        cros_collections.GroupNamedtuplesByKey(input_iter, 'test'),
        expected_result)


class InvertDictionayTests(cros_test_lib.TestCase):
  """Tests for InvertDictionary."""

  def testInvertDictionary(self):
    """Test InvertDictionary."""
    changes = ['change_1', 'change_2', 'change_3', 'change_4']
    slaves = ['slave_1', 'slave_2', 'slave_3', 'slave_4']
    slave_changes_dict = {
        slaves[0]: set(changes[0:1]),
        slaves[1]: set(changes[0:2]),
        slaves[2]: set(changes[2:4]),
        slaves[3]: set()
    }
    change_slaves_dict = cros_collections.InvertDictionary(
        slave_changes_dict)

    expected_dict = {
        changes[0]: set(slaves[0:2]),
        changes[1]: set([slaves[1]]),
        changes[2]: set([slaves[2]]),
        changes[3]: set([slaves[2]])
    }
    self.assertDictEqual(change_slaves_dict, expected_dict)