def test__attach_balances_series_bad(self):
tree = TreeNode.read([u"((a,b)c,d)r;"])
balances = pd.Series([10, -10])
with self.assertRaises(KeyError):
_attach_balances(balances, tree)
def test__attach_balances_series(self):
tree = TreeNode.read([u"((a,b)c,d)r;"])
balances = pd.Series([10, -10], index=['r', 'c'])
res_tree = _attach_balances(balances, tree)
self.assertEqual(res_tree.weight, 10)
def test__attach_balances_level_order(self):
tree = TreeNode.read([u"((a,b)c,d)r;"])
balances = np.array([10, -10])
res_tree = _attach_balances(balances, tree)
self.assertEqual(res_tree.weight, 10)
self.assertEqual(res_tree.children[0].weight, -10)
def test__attach_balances_bad_index(self):
tree = TreeNode.read([u"((a,b)c,d)r;"])
balances = np.array([10])
with self.assertRaises(IndexError):
_attach_balances(balances, tree)
def test__attach_balances(self):
tree = TreeNode.read([u"(a,b);"])
balances = np.array([10])
res_tree = _attach_balances(balances, tree)
self.assertEqual(res_tree.weight, 10)