#!/usr/bin/env python2 # Copyright 2015 Google Inc. All Rights Reserved. """Unit tests for the MachineImageManager class.""" from __future__ import print_function import random import unittest from machine_image_manager import MachineImageManager class MockLabel(object): """Class for generating a mock Label.""" def __init__(self, name, remotes=None): self.name = name self.remote = remotes def __hash__(self): """Provide hash function for label. This is required because Label object is used inside a dict as key. """ return hash(self.name) def __eq__(self, other): """Provide eq function for label. This is required because Label object is used inside a dict as key. """ return isinstance(other, MockLabel) and other.name == self.name class MockDut(object): """Class for creating a mock Device-Under-Test (DUT).""" def __init__(self, name, label=None): self.name = name self.label_ = label class MachineImageManagerTester(unittest.TestCase): """Class for testing MachineImageManager.""" def gen_duts_by_name(self, *names): duts = [] for n in names: duts.append(MockDut(n)) return duts def print_matrix(self, matrix): for r in matrix: for v in r: print('{} '.format('.' if v == ' ' else v)), print('') def create_labels_and_duts_from_pattern(self, pattern): labels = [] duts = [] for i, r in enumerate(pattern): l = MockLabel('l{}'.format(i), []) for j, v in enumerate(r.split()): if v == '.': l.remote.append('m{}'.format(j)) if i == 0: duts.append(MockDut('m{}'.format(j))) labels.append(l) return labels, duts def check_matrix_against_pattern(self, matrix, pattern): for i, s in enumerate(pattern): for j, v in enumerate(s.split()): self.assertTrue(v == '.' and matrix[i][j] == ' ' or v == matrix[i][j]) def pattern_based_test(self, inp, output): labels, duts = self.create_labels_and_duts_from_pattern(inp) mim = MachineImageManager(labels, duts) self.assertTrue(mim.compute_initial_allocation()) self.check_matrix_against_pattern(mim.matrix_, output) return mim def test_single_dut(self): labels = [MockLabel('l1'), MockLabel('l2'), MockLabel('l3')] dut = MockDut('m1') mim = MachineImageManager(labels, [dut]) mim.compute_initial_allocation() self.assertTrue(mim.matrix_ == [['Y'], ['Y'], ['Y']]) def test_single_label(self): labels = [MockLabel('l1')] duts = self.gen_duts_by_name('m1', 'm2', 'm3') mim = MachineImageManager(labels, duts) mim.compute_initial_allocation() self.assertTrue(mim.matrix_ == [['Y', 'Y', 'Y']]) def test_case1(self): labels = [MockLabel('l1', ['m1', 'm2']), MockLabel('l2', ['m2', 'm3']), MockLabel('l3', ['m1'])] duts = [MockDut('m1'), MockDut('m2'), MockDut('m3')] mim = MachineImageManager(labels, duts) self.assertTrue(mim.matrix_ == [[' ', ' ', 'X'], ['X', ' ', ' '], [' ', 'X', 'X']]) mim.compute_initial_allocation() self.assertTrue(mim.matrix_ == [[' ', 'Y', 'X'], ['X', ' ', 'Y'], ['Y', 'X', 'X']]) def test_case2(self): labels = [MockLabel('l1', ['m1', 'm2']), MockLabel('l2', ['m2', 'm3']), MockLabel('l3', ['m1'])] duts = [MockDut('m1'), MockDut('m2'), MockDut('m3')] mim = MachineImageManager(labels, duts) self.assertTrue(mim.matrix_ == [[' ', ' ', 'X'], ['X', ' ', ' '], [' ', 'X', 'X']]) mim.compute_initial_allocation() self.assertTrue(mim.matrix_ == [[' ', 'Y', 'X'], ['X', ' ', 'Y'], ['Y', 'X', 'X']]) def test_case3(self): labels = [MockLabel('l1', ['m1', 'm2']), MockLabel('l2', ['m2', 'm3']), MockLabel('l3', ['m1'])] duts = [MockDut('m1', labels[0]), MockDut('m2'), MockDut('m3')] mim = MachineImageManager(labels, duts) mim.compute_initial_allocation() self.assertTrue(mim.matrix_ == [[' ', 'Y', 'X'], ['X', ' ', 'Y'], ['Y', 'X', 'X']]) def test_case4(self): labels = [MockLabel('l1', ['m1', 'm2']), MockLabel('l2', ['m2', 'm3']), MockLabel('l3', ['m1'])] duts = [MockDut('m1'), MockDut('m2', labels[0]), MockDut('m3')] mim = MachineImageManager(labels, duts) mim.compute_initial_allocation() self.assertTrue(mim.matrix_ == [[' ', 'Y', 'X'], ['X', ' ', 'Y'], ['Y', 'X', 'X']]) def test_case5(self): labels = [MockLabel('l1', ['m3']), MockLabel('l2', ['m3']), MockLabel('l3', ['m1'])] duts = self.gen_duts_by_name('m1', 'm2', 'm3') mim = MachineImageManager(labels, duts) self.assertTrue(mim.compute_initial_allocation()) self.assertTrue(mim.matrix_ == [['X', 'X', 'Y'], ['X', 'X', 'Y'], ['Y', 'X', 'X']]) def test_2x2_with_allocation(self): labels = [MockLabel('l0'), MockLabel('l1')] duts = [MockDut('m0'), MockDut('m1')] mim = MachineImageManager(labels, duts) self.assertTrue(mim.compute_initial_allocation()) self.assertTrue(mim.allocate(duts[0]) == labels[0]) self.assertTrue(mim.allocate(duts[0]) == labels[1]) self.assertTrue(mim.allocate(duts[0]) is None) self.assertTrue(mim.matrix_[0][0] == '_') self.assertTrue(mim.matrix_[1][0] == '_') self.assertTrue(mim.allocate(duts[1]) == labels[1]) def test_10x10_general(self): """Gen 10x10 matrix.""" n = 10 labels = [] duts = [] for i in range(n): labels.append(MockLabel('l{}'.format(i))) duts.append(MockDut('m{}'.format(i))) mim = MachineImageManager(labels, duts) self.assertTrue(mim.compute_initial_allocation()) for i in range(n): for j in range(n): if i == j: self.assertTrue(mim.matrix_[i][j] == 'Y') else: self.assertTrue(mim.matrix_[i][j] == ' ') self.assertTrue(mim.allocate(duts[3]).name == 'l3') def test_random_generated(self): n = 10 labels = [] duts = [] for i in range(10): # generate 3-5 machines that is compatible with this label l = MockLabel('l{}'.format(i), []) r = random.random() for _ in range(4): t = int(r * 10) % n r *= 10 l.remote.append('m{}'.format(t)) labels.append(l) duts.append(MockDut('m{}'.format(i))) mim = MachineImageManager(labels, duts) self.assertTrue(mim.compute_initial_allocation()) def test_10x10_fully_random(self): inp = ['X . . . X X . X X .', 'X X . X . X . X X .', 'X X X . . X . X . X', 'X . X X . . X X . X', 'X X X X . . . X . .', 'X X . X . X . . X .', '. X . X . X X X . .', '. X . X X . X X . .', 'X X . . . X X X . .', '. X X X X . . . . X'] output = ['X Y . . X X . X X .', 'X X Y X . X . X X .', 'X X X Y . X . X . X', 'X . X X Y . X X . X', 'X X X X . Y . X . .', 'X X . X . X Y . X .', 'Y X . X . X X X . .', '. X . X X . X X Y .', 'X X . . . X X X . Y', '. X X X X . . Y . X'] self.pattern_based_test(inp, output) def test_10x10_fully_random2(self): inp = ['X . X . . X . X X X', 'X X X X X X . . X .', 'X . X X X X X . . X', 'X X X . X . X X . .', '. X . X . X X X X X', 'X X X X X X X . . X', 'X . X X X X X . . X', 'X X X . X X X X . .', 'X X X . . . X X X X', '. X X . X X X . X X'] output = ['X . X Y . X . X X X', 'X X X X X X Y . X .', 'X Y X X X X X . . X', 'X X X . X Y X X . .', '. X Y X . X X X X X', 'X X X X X X X Y . X', 'X . X X X X X . Y X', 'X X X . X X X X . Y', 'X X X . Y . X X X X', 'Y X X . X X X . X X'] self.pattern_based_test(inp, output) def test_3x4_with_allocation(self): inp = ['X X . .', '. . X .', 'X . X .'] output = ['X X Y .', 'Y . X .', 'X Y X .'] mim = self.pattern_based_test(inp, output) self.assertTrue(mim.allocate(mim.duts_[2]) == mim.labels_[0]) self.assertTrue(mim.allocate(mim.duts_[3]) == mim.labels_[2]) self.assertTrue(mim.allocate(mim.duts_[0]) == mim.labels_[1]) self.assertTrue(mim.allocate(mim.duts_[1]) == mim.labels_[2]) self.assertTrue(mim.allocate(mim.duts_[3]) == mim.labels_[1]) self.assertTrue(mim.allocate(mim.duts_[3]) == mim.labels_[0]) self.assertTrue(mim.allocate(mim.duts_[3]) is None) self.assertTrue(mim.allocate(mim.duts_[2]) is None) self.assertTrue(mim.allocate(mim.duts_[1]) == mim.labels_[1]) self.assertTrue(mim.allocate(mim.duts_[1]) == None) self.assertTrue(mim.allocate(mim.duts_[0]) == None) self.assertTrue(mim.label_duts_[0] == [2, 3]) self.assertTrue(mim.label_duts_[1] == [0, 3, 1]) self.assertTrue(mim.label_duts_[2] == [3, 1]) self.assertTrue(mim.allocate_log_ == [(0, 2), (2, 3), (1, 0), (2, 1), (1, 3), (0, 3), (1, 1)]) def test_cornercase_1(self): """This corner case is brought up by Caroline. The description is - If you have multiple labels and multiple machines, (so we don't automatically fall into the 1 dut or 1 label case), but all of the labels specify the same 1 remote, then instead of assigning the same machine to all the labels, your algorithm fails to assign any... So first step is to create an initial matrix like below, l0, l1 and l2 all specify the same 1 remote - m0. m0 m1 m2 l0 . X X l1 . X X l2 . X X The search process will be like this - a) try to find a solution with at most 1 'Y's per column (but ensure at least 1 Y per row), fail b) try to find a solution with at most 2 'Y's per column (but ensure at least 1 Y per row), fail c) try to find a solution with at most 3 'Y's per column (but ensure at least 1 Y per row), succeed, so we end up having this solution m0 m1 m2 l0 Y X X l1 Y X X l2 Y X X """ inp = ['. X X', '. X X', '. X X'] output = ['Y X X', 'Y X X', 'Y X X'] mim = self.pattern_based_test(inp, output) self.assertTrue(mim.allocate(mim.duts_[1]) is None) self.assertTrue(mim.allocate(mim.duts_[2]) is None) self.assertTrue(mim.allocate(mim.duts_[0]) == mim.labels_[0]) self.assertTrue(mim.allocate(mim.duts_[0]) == mim.labels_[1]) self.assertTrue(mim.allocate(mim.duts_[0]) == mim.labels_[2]) self.assertTrue(mim.allocate(mim.duts_[0]) is None) if __name__ == '__main__': unittest.main()