Question 3 - Linking nodes in a tree

#!/usr/bin/python

import unittest

class Node(object):
    def __init__(self, data, children=None):
        self.children = []
        self.right = None
        self.data = data

    def __str__(self):
        return str(self.data)

    def __repr__(self):
        return self.__str__()

def set_right_node(root):
    this = [root]
    last = None
    while this:
        next = []
        for i, node in enumerate(this):
            node.right = this[i+1] if i < len(this)-1 else None
            for i, c in enumerate(node.children):
                next.append(c)
        this = next

class Test(unittest.TestCase):
    def test_set_right_node(self):

        #      1
        #    / | \
        #   2  3  4
        #  / \     \
        # 5   6     7
        root_node = Node(1)
        node2 = Node(2)
        node3 = Node(3)
        node4 = Node(4)
        root_node.children = [node2, node3, node4]
        node5 = Node(5)
        node6 = Node(6)
        node2.children = [node5, node6]
        node7 = Node(7)
        node4.children = [node7]

        set_right_node(root_node)

        self.assertEqual(root_node.right, None)
        self.assertEqual(node2.right, node3)
        self.assertEqual(node3.right, node4)
        self.assertEqual(node4.right, None)
        self.assertEqual(node5.right, node6)
        self.assertEqual(node6.right, node7)
        self.assertEqual(node7.right, None)

if __name__ == '__main__':
    unittest.main()