nilmdb/tests/test_rbtree.py

# -*- coding: utf-8 -*-

import nilmdb
from nilmdb.utils.printf import *

from nose.tools import *
from nose.tools import assert_raises

from nilmdb.rbtree import RBTree, RBNode

from testutil.helpers import *
import unittest

# set to False to skip live renders
do_live_renders = False
def render(tree, description = "", live = True):
    import testutil.renderdot as renderdot
    r = renderdot.RBTreeRenderer(tree)
    return r.render(description, live and do_live_renders)

class TestRBTree:
    def test_rbtree(self):
        rb = RBTree()
        rb.insert(RBNode(10000, 10001))
        rb.insert(RBNode(10004, 10007))
        rb.insert(RBNode(10001, 10002))
        # There was a typo that gave the RBTree a loop in this case.
        # Verify that the dot isn't too big.
        s = render(rb, live = False)
        assert(len(s.splitlines()) < 30)

    def test_rbtree_big(self):
        import random
        random.seed(1234)

        # make a set of 100 intervals, inserted in order
        rb = RBTree()
        j = 100
        for i in xrange(j):
            rb.insert(RBNode(i, i+1))
        render(rb, "in-order insert")

        # remove about half of them
        for i in random.sample(xrange(j),j):
            if random.randint(0,1):
                rb.delete(rb.find(i, i+1))
        render(rb, "in-order insert, random delete")

        # make a set of 100 intervals, inserted at random
        rb = RBTree()
        j = 100
        for i in random.sample(xrange(j),j):
            rb.insert(RBNode(i, i+1))
        render(rb, "random insert")

        # remove about half of them
        for i in random.sample(xrange(j),j):
            if random.randint(0,1):
                rb.delete(rb.find(i, i+1))
        render(rb, "random insert, random delete")

        # in-order insert of 50 more
        for i in xrange(50):
            rb.insert(RBNode(i+500, i+501))
        render(rb, "random insert, random delete, in-order insert")

    def test_rbtree_basics(self):
        rb = RBTree()
        vals = [ 7, 14, 1, 2, 8, 11, 5, 15, 4]
        for n in vals:
            rb.insert(RBNode(n, n))

        # stringify
        s = ""
        for node in rb:
            s += str(node)
        in_("[node (None) 1", s)
        eq_(str(rb.nil), "[node nil]")

        # inorder traversal, successor and predecessor
        last = 0
        for node in rb:
            assert(node.start > last)
            last = node.start
            successor = rb.successor(node)
            if successor:
                assert(rb.predecessor(successor) is node)
            predecessor = rb.predecessor(node)
            if predecessor:
                assert(rb.successor(predecessor) is node)

        # Delete node not in the tree
        with assert_raises(AttributeError):
            rb.delete(RBNode(1,2))

        # Delete all nodes!
        for node in rb:
            rb.delete(node)

        # Build it up again, make sure it matches
        for n in vals:
            rb.insert(RBNode(n, n))
        s2 = ""
        for node in rb:
            s2 += str(node)
        assert(s == s2)

    def test_rbtree_find(self):
        # Get a little bit of coverage for some overlapping cases,
        # even though the class doesn't fully support it.
        rb = RBTree()
        nodes = [ RBNode(1, 5), RBNode(1, 10), RBNode(1, 15) ]
        for n in nodes:
            rb.insert(n)
        assert(rb.find(1, 5) is nodes[0])
        assert(rb.find(1, 10) is nodes[1])
        assert(rb.find(1, 15) is nodes[2])

    def test_rbtree_find_leftright(self):
        # Now let's get some ranges in there
        rb = RBTree()
        vals = [ 7, 14, 1, 2, 8, 11, 5, 15, 4]
        for n in vals:
            rb.insert(RBNode(n*10, n*10+5))

        # Check find_end_left, find_right_start
        for i in range(160):
            left = rb.find_left_end(i)
            right = rb.find_right_start(i)
            if left:
                # endpoint should be more than i
                assert(left.end >= i)
                # all earlier nodes should have a lower endpoint
                for node in rb:
                    if node is left:
                        break
                    assert(node.end < i)
            if right:
                # startpoint should be less than i
                assert(right.start <= i)
                # all later nodes should have a higher startpoint
                for node in reversed(list(rb)):
                    if node is right:
                        break
                    assert(node.start > i)

    def test_rbtree_intersect(self):
        # Fill with some ranges
        rb = RBTree()
        rb.insert(RBNode(10,20))
        rb.insert(RBNode(20,25))
        rb.insert(RBNode(30,40))
        # Just a quick test; test_interval will do better.
        eq_(len(list(rb.intersect(1,100))), 3)
        eq_(len(list(rb.intersect(10,20))), 1)
        eq_(len(list(rb.intersect(5,15))), 1)
        eq_(len(list(rb.intersect(15,15))), 1)
        eq_(len(list(rb.intersect(20,21))), 1)
        eq_(len(list(rb.intersect(19,21))), 2)