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