merkle-tree-stream/test/test_tree.py

65 lines
1.4 KiB
Python

"""Merkle tree test module."""
import hashlib
import pytest
from merkle_tree_stream import MerkleTreeIterator, MerkleTreeNode
def test_hashes(leaf, parent):
merkle_iter = MerkleTreeIterator(leaf=leaf, parent=parent)
merkle_iter.write(b'a')
merkle_iter.write(b'b')
expected_count = 2 + 1 # nodes plus parent
assert len(merkle_iter) == expected_count
assert next(merkle_iter) == MerkleTreeNode(
index=0,
parent=1,
hash=hashlib.sha256(b'a').hexdigest(),
size=1,
data=b'a',
)
assert next(merkle_iter) == MerkleTreeNode(
index=2,
parent=1,
hash=hashlib.sha256(b'b').hexdigest(),
size=1,
data=b'b',
)
hashed = hashlib.sha256(b'a')
hashed.update(b'b')
assert next(merkle_iter) == MerkleTreeNode(
index=1, parent=3, hash=hashed.hexdigest(), size=2, data=b''
)
with pytest.raises(StopIteration):
next(merkle_iter)
def test_single_root(leaf, parent):
merkle_iter = MerkleTreeIterator(leaf=leaf, parent=parent)
merkle_iter.write(b'a')
merkle_iter.write(b'b')
merkle_iter.write(b'c')
merkle_iter.write(b'd')
assert len(merkle_iter.roots) == 1
def multiple_roots(leaf, parent):
merkle_iter = MerkleTreeIterator(leaf=leaf, parent=parent)
merkle_iter.write(b'a')
merkle_iter.write(b'b')
merkle_iter.write(b'c')
assert len(merkle_iter.roots) > 1