| #!/usr/bin/python |
| |
| # Copyright 2015 Google Inc. All Rights Reserved. |
| |
| import random |
| import unittest |
| |
| from machine_image_manager import MachineImageManager |
| |
| class MockLabel(object): |
| |
| 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): |
| |
| def __init__(self, name, label=None): |
| self.name = name |
| self.label_ = label |
| |
| |
| class MachineImageManagerTester(unittest.TestCase): |
| |
| 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, input, output): |
| labels, duts = self.create_labels_and_duts_from_pattern(input) |
| 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): |
| input = ['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(input, output) |
| |
| def test_10x10_fully_random2(self): |
| input = ['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(input, output) |
| |
| def test_3x4_with_allocation(self): |
| input = ['X X . .', |
| '. . X .', |
| 'X . X .'] |
| output = ['X X Y .', |
| 'Y . X .', |
| 'X Y X .'] |
| mim = self.pattern_based_test(input, output) |
| self.assertTrue(mim.allocate(mim.duts_[2]) == mim.labels_[0]) |
| self.assertTrue(mim.allocate(mim.duts_[3]) == mim.labels_[1]) |
| 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_[0]) |
| self.assertTrue(mim.allocate(mim.duts_[3]) == mim.labels_[2]) |
| 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] == [3, 0, 1]) |
| self.assertTrue(mim.label_duts_[2] == [1, 3]) |
| self.assertTrue(mim.allocate_log_ == |
| [(0, 2), |
| (1, 3), |
| (1, 0), |
| (2, 1), |
| (0, 3), |
| (2, 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 |
| |
| """ |
| |
| input = ['. X X', |
| '. X X', |
| '. X X', ] |
| output = ['Y X X', |
| 'Y X X', |
| 'Y X X', ] |
| mim = self.pattern_based_test(input, 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() |