merkle-tree-stream/test/test_tree.py

57 lines
1.3 KiB
Python
Raw Normal View History

2019-08-07 05:58:26 +00:00
"""Merkle tree test module."""
import hashlib
import pytest
from merkle_tree_stream import MerkleTreeGenerator, MerkleTreeNode
def test_hashes(leaf, parent):
merkle = MerkleTreeGenerator(leaf=leaf, parent=parent)
merkle.write(b'a')
merkle.write(b'b')
expected_count = 2 + 1 # nodes plus parent
assert len(merkle) == expected_count
assert next(merkle) == MerkleTreeNode(
index=0, parent=1, hash=hashlib.sha256(b'a').digest(), size=1, data=b'a'
)
assert next(merkle) == 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) == MerkleTreeNode(
index=1, parent=3, hash=hashed.digest(), size=2, data=b''
)
with pytest.raises(StopIteration):
next(merkle)
def test_single_root(leaf, parent):
merkle = MerkleTreeGenerator(leaf=leaf, parent=parent)
merkle.write(b'a')
merkle.write(b'b')
merkle.write(b'c')
merkle.write(b'd')
assert len(merkle.roots) == 1
def multiple_roots(leaf, parent):
merkle = MerkleTreeGenerator(leaf=leaf, parent=parent)
merkle.write(b'a')
merkle.write(b'b')
merkle.write(b'c')
assert len(merkle.roots) > 1