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