merkle-tree-stream/test/test_tree.py

"""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