machine_image_manager_unittest.py revision e1a28bd82a8718b67a370433dd41932c567f6899 
1#!/usr/bin/python2
2
3# Copyright 2015 Google Inc. All Rights Reserved.
4
5"""Unit tests for the MachineImageManager class."""
6
7from __future__ import print_function
8
9import random
10import unittest
11
12from machine_image_manager import MachineImageManager
13
14
15class 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
37class 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
45class 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
289if __name__ == '__main__':
290  unittest.main()
291