Files

84 lines
2.3 KiB
Python

import random
from interval_tree import IntervalTree
def test_empty():
t = IntervalTree()
assert t.query_point(5) == []
assert len(t) == 0
def test_basic():
t = IntervalTree([(1, 5), (2, 3), (4, 8), (10, 12)])
assert sorted(t.query_point(3)) == [(1, 5), (2, 3)]
assert sorted(t.query_point(4)) == [(1, 5), (4, 8)]
assert sorted(t.query_point(11)) == [(10, 12)]
assert t.query_point(0) == []
assert t.query_point(100) == []
def test_point_intervals():
t = IntervalTree([(5, 5), (5, 5), (1, 9)])
assert sorted(t.query_point(5)) == [(1, 9), (5, 5), (5, 5)]
def test_insert_delete():
t = IntervalTree()
t.insert(1, 4)
t.insert(3, 6)
assert sorted(t.query_point(3)) == [(1, 4), (3, 6)]
assert t.delete(1, 4) is True
assert t.delete(1, 4) is False
assert t.query_point(3) == [(3, 6)]
assert len(t) == 1
def test_invalid():
try:
IntervalTree([(2, 1)])
raise AssertionError("expected ValueError")
except ValueError:
pass
t = IntervalTree()
try:
t.insert(5, 1)
raise AssertionError("expected ValueError")
except ValueError:
pass
def test_random_brute():
rng = random.Random(42)
for trial in range(200):
n = rng.randint(0, 40)
ivs = []
for _ in range(n):
a = rng.randint(-20, 20)
b = rng.randint(-20, 20)
lo, hi = (a, b) if a <= b else (b, a)
ivs.append((lo, hi))
t = IntervalTree(ivs)
for x in range(-25, 26):
got = sorted(t.query_point(x))
exp = sorted(iv for iv in ivs if iv[0] <= x <= iv[1])
assert got == exp, (trial, x, got, exp)
# mutate
for _ in range(5):
if rng.random() < 0.5 and ivs:
iv = rng.choice(ivs)
t.delete(*iv)
ivs.remove(iv)
else:
a = rng.randint(-20, 20); b = rng.randint(-20, 20)
lo, hi = (a, b) if a <= b else (b, a)
t.insert(lo, hi); ivs.append((lo, hi))
for x in range(-25, 26):
got = sorted(t.query_point(x))
exp = sorted(iv for iv in ivs if iv[0] <= x <= iv[1])
assert got == exp, (trial, x, got, exp)
if __name__ == "__main__":
import pytest, sys
sys.exit(pytest.main([__file__, "-q"]))