#!/usr/bin/env python2
# Copyright 2015 Google Inc. All Rights Reserved.
"""Unit tests for the MachineImageManager class."""
from __future__ import print_function
import random
import unittest
from machine_image_manager import MachineImageManager
class MockLabel(object):
"""Class for generating a mock Label."""
def __init__(self, name, remotes=None):
self.name = name
self.remote = remotes
def __hash__(self):
"""Provide hash function for label.
This is required because Label object is used inside a dict as key.
"""
return hash(self.name)
def __eq__(self, other):
"""Provide eq function for label.
This is required because Label object is used inside a dict as key.
"""
return isinstance(other, MockLabel) and other.name == self.name
class MockDut(object):
"""Class for creating a mock Device-Under-Test (DUT)."""
def __init__(self, name, label=None):
self.name = name
self.label_ = label
class MachineImageManagerTester(unittest.TestCase):
"""Class for testing MachineImageManager."""
def gen_duts_by_name(self, *names):
duts = []
for n in names:
duts.append(MockDut(n))
return duts
def print_matrix(self, matrix):
# pylint: disable=expression-not-assigned
for r in matrix:
for v in r:
print('{} '.format('.' if v == ' ' else v)),
print('')
def create_labels_and_duts_from_pattern(self, pattern):
labels = []
duts = []
for i, r in enumerate(pattern):
l = MockLabel('l{}'.format(i), [])
for j, v in enumerate(r.split()):
if v == '.':
l.remote.append('m{}'.format(j))
if i == 0:
duts.append(MockDut('m{}'.format(j)))
labels.append(l)
return labels, duts
def check_matrix_against_pattern(self, matrix, pattern):
for i, s in enumerate(pattern):
for j, v in enumerate(s.split()):
self.assertTrue(v == '.' and matrix[i][j] == ' ' or v == matrix[i][j])
def pattern_based_test(self, inp, output):
labels, duts = self.create_labels_and_duts_from_pattern(inp)
mim = MachineImageManager(labels, duts)
self.assertTrue(mim.compute_initial_allocation())
self.check_matrix_against_pattern(mim.matrix_, output)
return mim
def test_single_dut(self):
labels = [MockLabel('l1'), MockLabel('l2'), MockLabel('l3')]
dut = MockDut('m1')
mim = MachineImageManager(labels, [dut])
mim.compute_initial_allocation()
self.assertTrue(mim.matrix_ == [['Y'], ['Y'], ['Y']])
def test_single_label(self):
labels = [MockLabel('l1')]
duts = self.gen_duts_by_name('m1', 'm2', 'm3')
mim = MachineImageManager(labels, duts)
mim.compute_initial_allocation()
self.assertTrue(mim.matrix_ == [['Y', 'Y', 'Y']])
def test_case1(self):
labels = [
MockLabel('l1', ['m1', 'm2']), MockLabel('l2', ['m2', 'm3']), MockLabel(
'l3', ['m1'])
]
duts = [MockDut('m1'), MockDut('m2'), MockDut('m3')]
mim = MachineImageManager(labels, duts)
self.assertTrue(mim.matrix_ == [[' ', ' ', 'X'], ['X', ' ', ' '],
[' ', 'X', 'X']])
mim.compute_initial_allocation()
self.assertTrue(mim.matrix_ == [[' ', 'Y', 'X'], ['X', ' ', 'Y'],
['Y', 'X', 'X']])
def test_case2(self):
labels = [
MockLabel('l1', ['m1', 'm2']), MockLabel('l2', ['m2', 'm3']), MockLabel(
'l3', ['m1'])
]
duts = [MockDut('m1'), MockDut('m2'), MockDut('m3')]
mim = MachineImageManager(labels, duts)
self.assertTrue(mim.matrix_ == [[' ', ' ', 'X'], ['X', ' ', ' '],
[' ', 'X', 'X']])
mim.compute_initial_allocation()
self.assertTrue(mim.matrix_ == [[' ', 'Y', 'X'], ['X', ' ', 'Y'],
['Y', 'X', 'X']])
def test_case3(self):
labels = [
MockLabel('l1', ['m1', 'm2']), MockLabel('l2', ['m2', 'm3']), MockLabel(
'l3', ['m1'])
]
duts = [MockDut('m1', labels[0]), MockDut('m2'), MockDut('m3')]
mim = MachineImageManager(labels, duts)
mim.compute_initial_allocation()
self.assertTrue(mim.matrix_ == [[' ', 'Y', 'X'], ['X', ' ', 'Y'],
['Y', 'X', 'X']])
def test_case4(self):
labels = [
MockLabel('l1', ['m1', 'm2']), MockLabel('l2', ['m2', 'm3']), MockLabel(
'l3', ['m1'])
]
duts = [MockDut('m1'), MockDut('m2', labels[0]), MockDut('m3')]
mim = MachineImageManager(labels, duts)
mim.compute_initial_allocation()
self.assertTrue(mim.matrix_ == [[' ', 'Y', 'X'], ['X', ' ', 'Y'],
['Y', 'X', 'X']])
def test_case5(self):
labels = [
MockLabel('l1', ['m3']), MockLabel('l2', ['m3']), MockLabel(
'l3', ['m1'])
]
duts = self.gen_duts_by_name('m1', 'm2', 'm3')
mim = MachineImageManager(labels, duts)
self.assertTrue(mim.compute_initial_allocation())
self.assertTrue(mim.matrix_ == [['X', 'X', 'Y'], ['X', 'X', 'Y'],
['Y', 'X', 'X']])
def test_2x2_with_allocation(self):
labels = [MockLabel('l0'), MockLabel('l1')]
duts = [MockDut('m0'), MockDut('m1')]
mim = MachineImageManager(labels, duts)
self.assertTrue(mim.compute_initial_allocation())
self.assertTrue(mim.allocate(duts[0]) == labels[0])
self.assertTrue(mim.allocate(duts[0]) == labels[1])
self.assertTrue(mim.allocate(duts[0]) is None)
self.assertTrue(mim.matrix_[0][0] == '_')
self.assertTrue(mim.matrix_[1][0] == '_')
self.assertTrue(mim.allocate(duts[1]) == labels[1])
def test_10x10_general(self):
"""Gen 10x10 matrix."""
n = 10
labels = []
duts = []
for i in range(n):
labels.append(MockLabel('l{}'.format(i)))
duts.append(MockDut('m{}'.format(i)))
mim = MachineImageManager(labels, duts)
self.assertTrue(mim.compute_initial_allocation())
for i in range(n):
for j in range(n):
if i == j:
self.assertTrue(mim.matrix_[i][j] == 'Y')
else:
self.assertTrue(mim.matrix_[i][j] == ' ')
self.assertTrue(mim.allocate(duts[3]).name == 'l3')
def test_random_generated(self):
n = 10
labels = []
duts = []
for i in range(10):
# generate 3-5 machines that is compatible with this label
l = MockLabel('l{}'.format(i), [])
r = random.random()
for _ in range(4):
t = int(r * 10) % n
r *= 10
l.remote.append('m{}'.format(t))
labels.append(l)
duts.append(MockDut('m{}'.format(i)))
mim = MachineImageManager(labels, duts)
self.assertTrue(mim.compute_initial_allocation())
def test_10x10_fully_random(self):
inp = [
'X . . . X X . X X .', 'X X . X . X . X X .',
'X X X . . X . X . X', 'X . X X . . X X . X',
'X X X X . . . X . .', 'X X . X . X . . X .',
'. X . X . X X X . .', '. X . X X . X X . .',
'X X . . . X X X . .', '. X X X X . . . . X'
]
output = [
'X Y . . X X . X X .', 'X X Y X . X . X X .',
'X X X Y . X . X . X', 'X . X X Y . X X . X',
'X X X X . Y . X . .', 'X X . X . X Y . X .',
'Y X . X . X X X . .', '. X . X X . X X Y .',
'X X . . . X X X . Y', '. X X X X . . Y . X'
]
self.pattern_based_test(inp, output)
def test_10x10_fully_random2(self):
inp = [
'X . X . . X . X X X', 'X X X X X X . . X .',
'X . X X X X X . . X', 'X X X . X . X X . .',
'. X . X . X X X X X', 'X X X X X X X . . X',
'X . X X X X X . . X', 'X X X . X X X X . .',
'X X X . . . X X X X', '. X X . X X X . X X'
]
output = [
'X . X Y . X . X X X', 'X X X X X X Y . X .',
'X Y X X X X X . . X', 'X X X . X Y X X . .',
'. X Y X . X X X X X', 'X X X X X X X Y . X',
'X . X X X X X . Y X', 'X X X . X X X X . Y',
'X X X . Y . X X X X', 'Y X X . X X X . X X'
]
self.pattern_based_test(inp, output)
def test_3x4_with_allocation(self):
inp = ['X X . .', '. . X .', 'X . X .']
output = ['X X Y .', 'Y . X .', 'X Y X .']
mim = self.pattern_based_test(inp, output)
self.assertTrue(mim.allocate(mim.duts_[2]) == mim.labels_[0])
self.assertTrue(mim.allocate(mim.duts_[3]) == mim.labels_[2])
self.assertTrue(mim.allocate(mim.duts_[0]) == mim.labels_[1])
self.assertTrue(mim.allocate(mim.duts_[1]) == mim.labels_[2])
self.assertTrue(mim.allocate(mim.duts_[3]) == mim.labels_[1])
self.assertTrue(mim.allocate(mim.duts_[3]) == mim.labels_[0])
self.assertTrue(mim.allocate(mim.duts_[3]) is None)
self.assertTrue(mim.allocate(mim.duts_[2]) is None)
self.assertTrue(mim.allocate(mim.duts_[1]) == mim.labels_[1])
self.assertTrue(mim.allocate(mim.duts_[1]) == None)
self.assertTrue(mim.allocate(mim.duts_[0]) == None)
self.assertTrue(mim.label_duts_[0] == [2, 3])
self.assertTrue(mim.label_duts_[1] == [0, 3, 1])
self.assertTrue(mim.label_duts_[2] == [3, 1])
self.assertTrue(mim.allocate_log_ == [(0, 2), (2, 3), (1, 0), (2, 1),
(1, 3), (0, 3), (1, 1)])
def test_cornercase_1(self):
"""This corner case is brought up by Caroline.
The description is -
If you have multiple labels and multiple machines, (so we don't
automatically fall into the 1 dut or 1 label case), but all of the
labels specify the same 1 remote, then instead of assigning the same
machine to all the labels, your algorithm fails to assign any...
So first step is to create an initial matrix like below, l0, l1 and l2
all specify the same 1 remote - m0.
m0 m1 m2
l0 . X X
l1 . X X
l2 . X X
The search process will be like this -
a) try to find a solution with at most 1 'Y's per column (but ensure at
least 1 Y per row), fail
b) try to find a solution with at most 2 'Y's per column (but ensure at
least 1 Y per row), fail
c) try to find a solution with at most 3 'Y's per column (but ensure at
least 1 Y per row), succeed, so we end up having this solution
m0 m1 m2
l0 Y X X
l1 Y X X
l2 Y X X
"""
inp = ['. X X', '. X X', '. X X']
output = ['Y X X', 'Y X X', 'Y X X']
mim = self.pattern_based_test(inp, output)
self.assertTrue(mim.allocate(mim.duts_[1]) is None)
self.assertTrue(mim.allocate(mim.duts_[2]) is None)
self.assertTrue(mim.allocate(mim.duts_[0]) == mim.labels_[0])
self.assertTrue(mim.allocate(mim.duts_[0]) == mim.labels_[1])
self.assertTrue(mim.allocate(mim.duts_[0]) == mim.labels_[2])
self.assertTrue(mim.allocate(mim.duts_[0]) is None)
if __name__ == '__main__':
unittest.main()