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