import unittest from ZTUtils import Tree class Item: children = () id = '' def __init__(self, id, children=()): self.id = id self.children = children def tpId(self): return self.id def tpValues(self): return self.children class TreeTests(unittest.TestCase): def setUp(self): self.tm = Tree.TreeMaker() self.root = Item('a', ( Item('b', ( Item('d'), Item('e'))), Item('c', ( Item('f', ( Item('h'), Item('i'))), Item('g'))))) self.items = { 'a': self.root, 'b': self.root.children[0], 'c': self.root.children[1], 'd': self.root.children[0].children[0], 'e': self.root.children[0].children[1], 'f': self.root.children[1].children[0], 'g': self.root.children[1].children[1], 'h': self.root.children[1].children[0].children[0], 'i': self.root.children[1].children[0].children[1], } self.expansionmap = {Tree.b2a('a'): {Tree.b2a('c'): None}} def testBaseTree(self): treeroot = self.tm.tree(self.root) self.assertEqual(len(treeroot), 2) self.assertEqual(treeroot.size, 3) self.assertEqual(treeroot.height, 2) self.assertEqual(treeroot.depth, 0) self.assertEqual(treeroot.state, 1) self.assert_(treeroot.object is self.root) i = 'b' for subnode in treeroot: self.assertEqual(len(subnode), 0) self.assertEqual(subnode.size, 1) self.assertEqual(subnode.height, 1) self.assertEqual(subnode.depth, 1) self.assertEqual(subnode.state, -1) self.assert_(subnode.object is self.items[i]) i = chr(ord(i) + 1) expected_set = [self.items['a'], self.items['b'], self.items['c']] set = treeroot.flat() self.assertEqual(len(set), treeroot.size) self.assertEqual([s.object for s in set], expected_set) set = [] def collect(node, set=set): set.append(node.object) treeroot.walk(collect) self.assertEqual(len(set), treeroot.size) self.assertEqual(set, expected_set) def testExpandedTree(self): treeroot = self.tm.tree(self.root, 1) self.assertEqual(len(treeroot), 2) self.assertEqual(treeroot.size, len(self.items)) self.assertEqual(treeroot.height, 4) self.assertEqual(treeroot.depth, 0) self.assertEqual(treeroot.state, 1) self.assert_(treeroot.object is self.root) items = self.items expected_set = [items['a'], items['b'], items['d'], items['e'], items['c'], items['f'], items['h'], items['i'], items['g']] set = treeroot.flat() self.assertEqual(len(set), treeroot.size) self.assertEqual([s.object for s in set], expected_set) set = [] def collect(node, set=set): set.append(node.object) treeroot.walk(collect) self.assertEqual(len(set), treeroot.size) self.assertEqual(set, expected_set) leaves = ('d', 'e', 'g', 'h', 'i') for node in treeroot.flat(): if node.object.id in leaves: self.assertEqual(node.state, 0) else: self.assertEqual(node.state, 1) self.assertEqual(node.size, len(node.flat())) def testExpansionMap(self): treeroot = self.tm.tree(self.root, self.expansionmap) self.assertEqual(treeroot.size, 5) self.assertEqual(treeroot.state, 1) self.assertEqual(treeroot[0].state, -1) self.assertEqual(treeroot[1].state, 1) self.assertEqual(treeroot[1][0].state, -1) self.assertEqual(treeroot[1][1].state, 0) def testAssumeChildren(self): self.tm.setAssumeChildren(True) treeroot = self.tm.tree(self.root, self.expansionmap) self.assertEqual(treeroot[1][1].state, -1) def testNoExpandRoot(self): self.tm.setExpandRoot(False) treeroot = self.tm.tree(self.root) self.assertEqual(treeroot.state, -1) self.assertEqual(len(treeroot), 0) def testIdAttribute(self): treeroot = self.tm.tree(self.root) self.tm.setIdAttr('id') treeroot2 = self.tm.tree(self.root) self.assertEqual(treeroot.id, treeroot2.id) def testChildrenAttribute(self): self.tm.setChildAccess(attrname='children') treeroot = self.tm.tree(self.root) self.assertEqual(len(treeroot), 2) def testChildrenFilter(self): def filter(children): return [c for c in children if c.id in ('b', 'd')] self.tm.setChildAccess(filter=filter) treeroot = self.tm.tree(self.root, 1) self.assertEqual(treeroot.size, 3) self.assertEqual(len(treeroot), 1) self.assertEqual(len(treeroot[0]), 1) expected_set = [self.items['a'], self.items['b'], self.items['d']] set = [] def collect(node, set=set): set.append(node.object) treeroot.walk(collect) self.assertEqual(set, expected_set) def testChildrenFunction(self): def childrenFunction(object): return object.children self.tm.setChildAccess(function=childrenFunction) treeroot = self.tm.tree(self.root) self.assertEqual(len(treeroot), 2) def testStateFunction(self): def stateFunction(object, state): if object.id == 'b': return 1 if object.id == 'd': return -1 return state self.tm.setStateFunction(stateFunction) treeroot = self.tm.tree(self.root) self.assertEqual(treeroot.size, 5) self.assertEqual(treeroot.state, 1) self.assertEqual(treeroot[0].state, 1) self.assertEqual(treeroot[0][0].state, -1) self.assertEqual(treeroot[0][1].state, 0) self.assertEqual(treeroot[1].state, -1) def testEncodeDecode(self): treeroot1 = self.tm.tree(self.root, self.expansionmap) encoded = Tree.encodeExpansion(treeroot1.flat()) self.failIf(encoded.find('\n') != -1) decodedmap = Tree.decodeExpansion(encoded) treeroot2 = self.tm.tree(self.root, decodedmap) self.assertEqual(treeroot1.size, treeroot2.size) self.assertEqual(len(treeroot1), len(treeroot2)) def testEncodedExpansionIdWithDot(self): # Regression test for Collector issue #603 # An encoded node ID with a first character with the first 6 bits set. item = Item('\xfcberbug!', (Item('b'),)) # 'uberbug!' with u-umlaut. treeroot1 = self.tm.tree(item) encoded = Tree.encodeExpansion(treeroot1.flat()) decodedmap = Tree.decodeExpansion(encoded) treeroot2 = self.tm.tree(item, decodedmap) self.assertEqual(treeroot1.size, treeroot2.size) self.assertEqual(len(treeroot1), len(treeroot2)) def testDecodeInputSizeLimit(self): self.assertRaises(ValueError, Tree.decodeExpansion, 'x' * 10000) def testDecodeDecompressedSizeLimit(self): import zlib from ZTUtils.Tree import b2a, a2b, encodeExpansion, decodeExpansion big = b2a(zlib.compress('x' * (1024*1100))) self.assert_(len(big) < 8192) # Must be under the input size limit self.assertRaises(ValueError, Tree.decodeExpansion, ':' + big) def test_suite(): return unittest.makeSuite(TreeTests) if __name__=='__main__': unittest.main()