TreeBuilder fix for leafs, test enhanced