57 lines
1.3 KiB
Python
57 lines
1.3 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').digest(), size=1, data=b'a'
|
|
)
|
|
|
|
assert next(merkle_iter) == MerkleTreeNode(
|
|
index=2, parent=1, hash=hashlib.sha256(b'b').digest(), 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.digest(), 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
|