You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

test_rbtree.py 5.0 KiB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159
  1. # -*- coding: utf-8 -*-
  2. import nilmdb
  3. from nilmdb.utils.printf import *
  4. from nose.tools import *
  5. from nose.tools import assert_raises
  6. from nilmdb.server.rbtree import RBTree, RBNode
  7. from testutil.helpers import *
  8. import unittest
  9. # set to False to skip live renders
  10. do_live_renders = False
  11. def render(tree, description = "", live = True):
  12. import testutil.renderdot as renderdot
  13. r = renderdot.RBTreeRenderer(tree)
  14. return r.render(description, live and do_live_renders)
  15. class TestRBTree:
  16. def test_rbtree(self):
  17. rb = RBTree()
  18. rb.insert(RBNode(10000, 10001))
  19. rb.insert(RBNode(10004, 10007))
  20. rb.insert(RBNode(10001, 10002))
  21. # There was a typo that gave the RBTree a loop in this case.
  22. # Verify that the dot isn't too big.
  23. s = render(rb, live = False)
  24. assert(len(s.splitlines()) < 30)
  25. def test_rbtree_big(self):
  26. import random
  27. random.seed(1234)
  28. # make a set of 100 intervals, inserted in order
  29. rb = RBTree()
  30. j = 100
  31. for i in range(j):
  32. rb.insert(RBNode(i, i+1))
  33. render(rb, "in-order insert")
  34. # remove about half of them
  35. for i in random.sample(range(j),j):
  36. if random.randint(0,1):
  37. rb.delete(rb.find(i, i+1))
  38. render(rb, "in-order insert, random delete")
  39. # make a set of 100 intervals, inserted at random
  40. rb = RBTree()
  41. j = 100
  42. for i in random.sample(range(j),j):
  43. rb.insert(RBNode(i, i+1))
  44. render(rb, "random insert")
  45. # remove about half of them
  46. for i in random.sample(range(j),j):
  47. if random.randint(0,1):
  48. rb.delete(rb.find(i, i+1))
  49. render(rb, "random insert, random delete")
  50. # in-order insert of 50 more
  51. for i in range(50):
  52. rb.insert(RBNode(i+500, i+501))
  53. render(rb, "random insert, random delete, in-order insert")
  54. def test_rbtree_basics(self):
  55. rb = RBTree()
  56. vals = [ 7, 14, 1, 2, 8, 11, 5, 15, 4]
  57. for n in vals:
  58. rb.insert(RBNode(n, n))
  59. # stringify
  60. s = ""
  61. for node in rb:
  62. s += str(node)
  63. in_("[node (None) 1", s)
  64. eq_(str(rb.nil), "[node nil]")
  65. # inorder traversal, successor and predecessor
  66. last = 0
  67. for node in rb:
  68. assert(node.start > last)
  69. last = node.start
  70. successor = rb.successor(node)
  71. if successor:
  72. assert(rb.predecessor(successor) is node)
  73. predecessor = rb.predecessor(node)
  74. if predecessor:
  75. assert(rb.successor(predecessor) is node)
  76. # Delete node not in the tree
  77. with assert_raises(AttributeError):
  78. rb.delete(RBNode(1,2))
  79. # Delete all nodes!
  80. for node in rb:
  81. rb.delete(node)
  82. # Build it up again, make sure it matches
  83. for n in vals:
  84. rb.insert(RBNode(n, n))
  85. s2 = ""
  86. for node in rb:
  87. s2 += str(node)
  88. assert(s == s2)
  89. def test_rbtree_find(self):
  90. # Get a little bit of coverage for some overlapping cases,
  91. # even though the class doesn't fully support it.
  92. rb = RBTree()
  93. nodes = [ RBNode(1, 5), RBNode(1, 10), RBNode(1, 15) ]
  94. for n in nodes:
  95. rb.insert(n)
  96. assert(rb.find(1, 5) is nodes[0])
  97. assert(rb.find(1, 10) is nodes[1])
  98. assert(rb.find(1, 15) is nodes[2])
  99. def test_rbtree_find_leftright(self):
  100. # Now let's get some ranges in there
  101. rb = RBTree()
  102. vals = [ 7, 14, 1, 2, 8, 11, 5, 15, 4]
  103. for n in vals:
  104. rb.insert(RBNode(n*10, n*10+5))
  105. # Check find_end_left, find_right_start
  106. for i in range(160):
  107. left = rb.find_left_end(i)
  108. right = rb.find_right_start(i)
  109. if left:
  110. # endpoint should be more than i
  111. assert(left.end >= i)
  112. # all earlier nodes should have a lower endpoint
  113. for node in rb:
  114. if node is left:
  115. break
  116. assert(node.end < i)
  117. if right:
  118. # startpoint should be less than i
  119. assert(right.start <= i)
  120. # all later nodes should have a higher startpoint
  121. for node in reversed(list(rb)):
  122. if node is right:
  123. break
  124. assert(node.start > i)
  125. def test_rbtree_intersect(self):
  126. # Fill with some ranges
  127. rb = RBTree()
  128. rb.insert(RBNode(10,20))
  129. rb.insert(RBNode(20,25))
  130. rb.insert(RBNode(30,40))
  131. # Just a quick test; test_interval will do better.
  132. eq_(len(list(rb.intersect(1,100))), 3)
  133. eq_(len(list(rb.intersect(10,20))), 1)
  134. eq_(len(list(rb.intersect(5,15))), 1)
  135. eq_(len(list(rb.intersect(15,15))), 1)
  136. eq_(len(list(rb.intersect(20,21))), 1)
  137. eq_(len(list(rb.intersect(19,21))), 2)