1 #!/usr/bin/env python2 2 3 # Copyright 2015 Google Inc. All Rights Reserved. 4 5 """Unit tests for the MachineImageManager class.""" 6 7 from __future__ import print_function 8 9 import random 10 import unittest 11 12 from machine_image_manager import MachineImageManager 13 14 15 class MockLabel(object): 16 """Class for generating a mock Label.""" 17 18 def __init__(self, name, remotes=None): 19 self.name = name 20 self.remote = remotes 21 22 def __hash__(self): 23 """Provide hash function for label. 24 25 This is required because Label object is used inside a dict as key. 26 """ 27 return hash(self.name) 28 29 def __eq__(self, other): 30 """Provide eq function for label. 31 32 This is required because Label object is used inside a dict as key. 33 """ 34 return isinstance(other, MockLabel) and other.name == self.name 35 36 37 class MockDut(object): 38 """Class for creating a mock Device-Under-Test (DUT).""" 39 40 def __init__(self, name, label=None): 41 self.name = name 42 self.label_ = label 43 44 45 class MachineImageManagerTester(unittest.TestCase): 46 """Class for testing MachineImageManager.""" 47 48 def gen_duts_by_name(self, *names): 49 duts = [] 50 for n in names: 51 duts.append(MockDut(n)) 52 return duts 53 54 def print_matrix(self, matrix): 55 for r in matrix: 56 for v in r: 57 print('{} '.format('.' if v == ' ' else v)), 58 print('') 59 60 def create_labels_and_duts_from_pattern(self, pattern): 61 labels = [] 62 duts = [] 63 for i, r in enumerate(pattern): 64 l = MockLabel('l{}'.format(i), []) 65 for j, v in enumerate(r.split()): 66 if v == '.': 67 l.remote.append('m{}'.format(j)) 68 if i == 0: 69 duts.append(MockDut('m{}'.format(j))) 70 labels.append(l) 71 return labels, duts 72 73 def check_matrix_against_pattern(self, matrix, pattern): 74 for i, s in enumerate(pattern): 75 for j, v in enumerate(s.split()): 76 self.assertTrue(v == '.' and matrix[i][j] == ' ' or v == matrix[i][j]) 77 78 def pattern_based_test(self, inp, output): 79 labels, duts = self.create_labels_and_duts_from_pattern(inp) 80 mim = MachineImageManager(labels, duts) 81 self.assertTrue(mim.compute_initial_allocation()) 82 self.check_matrix_against_pattern(mim.matrix_, output) 83 return mim 84 85 def test_single_dut(self): 86 labels = [MockLabel('l1'), MockLabel('l2'), MockLabel('l3')] 87 dut = MockDut('m1') 88 mim = MachineImageManager(labels, [dut]) 89 mim.compute_initial_allocation() 90 self.assertTrue(mim.matrix_ == [['Y'], ['Y'], ['Y']]) 91 92 def test_single_label(self): 93 labels = [MockLabel('l1')] 94 duts = self.gen_duts_by_name('m1', 'm2', 'm3') 95 mim = MachineImageManager(labels, duts) 96 mim.compute_initial_allocation() 97 self.assertTrue(mim.matrix_ == [['Y', 'Y', 'Y']]) 98 99 def test_case1(self): 100 labels = [MockLabel('l1', ['m1', 'm2']), MockLabel('l2', ['m2', 'm3']), 101 MockLabel('l3', ['m1'])] 102 duts = [MockDut('m1'), MockDut('m2'), MockDut('m3')] 103 mim = MachineImageManager(labels, duts) 104 self.assertTrue(mim.matrix_ == [[' ', ' ', 'X'], ['X', ' ', ' '], [' ', 'X', 105 'X']]) 106 mim.compute_initial_allocation() 107 self.assertTrue(mim.matrix_ == [[' ', 'Y', 'X'], ['X', ' ', 'Y'], ['Y', 'X', 108 'X']]) 109 110 def test_case2(self): 111 labels = [MockLabel('l1', ['m1', 'm2']), MockLabel('l2', ['m2', 'm3']), 112 MockLabel('l3', ['m1'])] 113 duts = [MockDut('m1'), MockDut('m2'), MockDut('m3')] 114 mim = MachineImageManager(labels, duts) 115 self.assertTrue(mim.matrix_ == [[' ', ' ', 'X'], ['X', ' ', ' '], [' ', 'X', 116 'X']]) 117 mim.compute_initial_allocation() 118 self.assertTrue(mim.matrix_ == [[' ', 'Y', 'X'], ['X', ' ', 'Y'], ['Y', 'X', 119 'X']]) 120 121 def test_case3(self): 122 labels = [MockLabel('l1', ['m1', 'm2']), MockLabel('l2', ['m2', 'm3']), 123 MockLabel('l3', ['m1'])] 124 duts = [MockDut('m1', labels[0]), MockDut('m2'), MockDut('m3')] 125 mim = MachineImageManager(labels, duts) 126 mim.compute_initial_allocation() 127 self.assertTrue(mim.matrix_ == [[' ', 'Y', 'X'], ['X', ' ', 'Y'], ['Y', 'X', 128 'X']]) 129 130 def test_case4(self): 131 labels = [MockLabel('l1', ['m1', 'm2']), MockLabel('l2', ['m2', 'm3']), 132 MockLabel('l3', ['m1'])] 133 duts = [MockDut('m1'), MockDut('m2', labels[0]), MockDut('m3')] 134 mim = MachineImageManager(labels, duts) 135 mim.compute_initial_allocation() 136 self.assertTrue(mim.matrix_ == [[' ', 'Y', 'X'], ['X', ' ', 'Y'], ['Y', 'X', 137 'X']]) 138 139 def test_case5(self): 140 labels = [MockLabel('l1', ['m3']), MockLabel('l2', ['m3']), 141 MockLabel('l3', ['m1'])] 142 duts = self.gen_duts_by_name('m1', 'm2', 'm3') 143 mim = MachineImageManager(labels, duts) 144 self.assertTrue(mim.compute_initial_allocation()) 145 self.assertTrue(mim.matrix_ == [['X', 'X', 'Y'], ['X', 'X', 'Y'], ['Y', 'X', 146 'X']]) 147 148 def test_2x2_with_allocation(self): 149 labels = [MockLabel('l0'), MockLabel('l1')] 150 duts = [MockDut('m0'), MockDut('m1')] 151 mim = MachineImageManager(labels, duts) 152 self.assertTrue(mim.compute_initial_allocation()) 153 self.assertTrue(mim.allocate(duts[0]) == labels[0]) 154 self.assertTrue(mim.allocate(duts[0]) == labels[1]) 155 self.assertTrue(mim.allocate(duts[0]) is None) 156 self.assertTrue(mim.matrix_[0][0] == '_') 157 self.assertTrue(mim.matrix_[1][0] == '_') 158 self.assertTrue(mim.allocate(duts[1]) == labels[1]) 159 160 def test_10x10_general(self): 161 """Gen 10x10 matrix.""" 162 n = 10 163 labels = [] 164 duts = [] 165 for i in range(n): 166 labels.append(MockLabel('l{}'.format(i))) 167 duts.append(MockDut('m{}'.format(i))) 168 mim = MachineImageManager(labels, duts) 169 self.assertTrue(mim.compute_initial_allocation()) 170 for i in range(n): 171 for j in range(n): 172 if i == j: 173 self.assertTrue(mim.matrix_[i][j] == 'Y') 174 else: 175 self.assertTrue(mim.matrix_[i][j] == ' ') 176 self.assertTrue(mim.allocate(duts[3]).name == 'l3') 177 178 def test_random_generated(self): 179 n = 10 180 labels = [] 181 duts = [] 182 for i in range(10): 183 # generate 3-5 machines that is compatible with this label 184 l = MockLabel('l{}'.format(i), []) 185 r = random.random() 186 for _ in range(4): 187 t = int(r * 10) % n 188 r *= 10 189 l.remote.append('m{}'.format(t)) 190 labels.append(l) 191 duts.append(MockDut('m{}'.format(i))) 192 mim = MachineImageManager(labels, duts) 193 self.assertTrue(mim.compute_initial_allocation()) 194 195 def test_10x10_fully_random(self): 196 inp = ['X . . . X X . X X .', 'X X . X . X . X X .', 197 'X X X . . X . X . X', 'X . X X . . X X . X', 198 'X X X X . . . X . .', 'X X . X . X . . X .', 199 '. X . X . X X X . .', '. X . X X . X X . .', 200 'X X . . . X X X . .', '. X X X X . . . . X'] 201 output = ['X Y . . X X . X X .', 'X X Y X . X . X X .', 202 'X X X Y . X . X . X', 'X . X X Y . X X . X', 203 'X X X X . Y . X . .', 'X X . X . X Y . X .', 204 'Y X . X . X X X . .', '. X . X X . X X Y .', 205 'X X . . . X X X . Y', '. X X X X . . Y . X'] 206 self.pattern_based_test(inp, output) 207 208 def test_10x10_fully_random2(self): 209 inp = ['X . X . . X . X X X', 'X X X X X X . . X .', 210 'X . X X X X X . . X', 'X X X . X . X X . .', 211 '. X . X . X X X X X', 'X X X X X X X . . X', 212 'X . X X X X X . . X', 'X X X . X X X X . .', 213 'X X X . . . X X X X', '. X X . X X X . X X'] 214 output = ['X . X Y . X . X X X', 'X X X X X X Y . X .', 215 'X Y X X X X X . . X', 'X X X . X Y X X . .', 216 '. X Y X . X X X X X', 'X X X X X X X Y . X', 217 'X . X X X X X . Y X', 'X X X . X X X X . Y', 218 'X X X . Y . X X X X', 'Y X X . X X X . X X'] 219 self.pattern_based_test(inp, output) 220 221 def test_3x4_with_allocation(self): 222 inp = ['X X . .', '. . X .', 'X . X .'] 223 output = ['X X Y .', 'Y . X .', 'X Y X .'] 224 mim = self.pattern_based_test(inp, output) 225 self.assertTrue(mim.allocate(mim.duts_[2]) == mim.labels_[0]) 226 self.assertTrue(mim.allocate(mim.duts_[3]) == mim.labels_[2]) 227 self.assertTrue(mim.allocate(mim.duts_[0]) == mim.labels_[1]) 228 self.assertTrue(mim.allocate(mim.duts_[1]) == mim.labels_[2]) 229 self.assertTrue(mim.allocate(mim.duts_[3]) == mim.labels_[1]) 230 self.assertTrue(mim.allocate(mim.duts_[3]) == mim.labels_[0]) 231 self.assertTrue(mim.allocate(mim.duts_[3]) is None) 232 self.assertTrue(mim.allocate(mim.duts_[2]) is None) 233 self.assertTrue(mim.allocate(mim.duts_[1]) == mim.labels_[1]) 234 self.assertTrue(mim.allocate(mim.duts_[1]) == None) 235 self.assertTrue(mim.allocate(mim.duts_[0]) == None) 236 self.assertTrue(mim.label_duts_[0] == [2, 3]) 237 self.assertTrue(mim.label_duts_[1] == [0, 3, 1]) 238 self.assertTrue(mim.label_duts_[2] == [3, 1]) 239 self.assertTrue(mim.allocate_log_ == [(0, 2), (2, 3), (1, 0), (2, 1), 240 (1, 3), (0, 3), (1, 1)]) 241 242 def test_cornercase_1(self): 243 """This corner case is brought up by Caroline. 244 245 The description is - 246 247 If you have multiple labels and multiple machines, (so we don't 248 automatically fall into the 1 dut or 1 label case), but all of the 249 labels specify the same 1 remote, then instead of assigning the same 250 machine to all the labels, your algorithm fails to assign any... 251 252 So first step is to create an initial matrix like below, l0, l1 and l2 253 all specify the same 1 remote - m0. 254 255 m0 m1 m2 256 l0 . X X 257 258 l1 . X X 259 260 l2 . X X 261 262 The search process will be like this - 263 a) try to find a solution with at most 1 'Y's per column (but ensure at 264 least 1 Y per row), fail 265 b) try to find a solution with at most 2 'Y's per column (but ensure at 266 least 1 Y per row), fail 267 c) try to find a solution with at most 3 'Y's per column (but ensure at 268 least 1 Y per row), succeed, so we end up having this solution 269 270 m0 m1 m2 271 l0 Y X X 272 273 l1 Y X X 274 275 l2 Y X X 276 """ 277 278 inp = ['. X X', '. X X', '. X X'] 279 output = ['Y X X', 'Y X X', 'Y X X'] 280 mim = self.pattern_based_test(inp, output) 281 self.assertTrue(mim.allocate(mim.duts_[1]) is None) 282 self.assertTrue(mim.allocate(mim.duts_[2]) is None) 283 self.assertTrue(mim.allocate(mim.duts_[0]) == mim.labels_[0]) 284 self.assertTrue(mim.allocate(mim.duts_[0]) == mim.labels_[1]) 285 self.assertTrue(mim.allocate(mim.duts_[0]) == mim.labels_[2]) 286 self.assertTrue(mim.allocate(mim.duts_[0]) is None) 287 288 289 if __name__ == '__main__': 290 unittest.main() 291
