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