def test_num_nodes():
t = Trie()
assert t.num_nodes() == 0
t[b"foo"] = 1
assert t.num_nodes() == 3
t[b"foobar"] = 1
assert t.num_nodes() == 6
t[b"foozle"] = 1
assert t.num_nodes() == 9
t[b"hello"] = 1
assert t.num_nodes() == 14
del t[b"foo"]
assert not b"foo" in t
assert t.num_nodes() == 14
del t[b"foozle"]
assert t.num_nodes() == 11
del t[b"foobar"]
assert t.num_nodes() == 5
del t[b"hello"]
assert t.num_nodes() == 0
n = 100
for i in xrange(n):
t[b(str(i))] = i
assert t.num_nodes() == n
def test_iter():
"""Test iterating over a Trie"""
t = Trie()
words = [b"hello", b"foo", b"foobar", b"foozle"]
for key in words:
t[key] = 1
assert set([b(key) for key in t]) == set(words)
# Modify values in the trie using t's iterator
for key in t:
t[b(key)] = 2
assert t.values() == [2]*len(words)
# It should be an error to continue to iterate after adding and/or removing
# nodes.
i1 = iter(t)
next(i1) # This should be fine
t[b"new"] = 1
with pytest.raises(RuntimeError):
next(i1)
# Alright, so addition was detected, now for deletion
i2 = iter(t)
del t[b"new"]
with pytest.raises(RuntimeError):
next(i2)
# Detection should not be on number of nodes, so let's add and remove a
# single node.
i3 = iter(t)
next(i3) # should be fine
n = t.num_nodes()
t[b"a"] = 1
assert t.num_nodes() == n + 1
del t[b"a"]
assert t.num_nodes() == n
with pytest.raises(RuntimeError):
next(i3)