test_traverse_tree.py

"""
TODO:
    - Rewrite to more clean code, not so repetitive
    - Update when signals and id's has been merged.
"""

from b_asic.core_operations import Constant, Addition
from b_asic.signal import Signal 
from b_asic.port import InputPort, OutputPort
from b_asic.traverse_tree import Traverse

import pytest

@pytest.fixture
def operation():
    return Constant(2)

def create_operation(_type, dest_oper, index, **kwargs):
    oper = _type(**kwargs)
    oper_signal = Signal()
    oper._output_ports[0].connect(oper_signal)

    dest_oper._input_ports[index].connect(oper_signal)
    return oper

@pytest.fixture
def operation_tree():
    add_oper = Addition()

    const_oper = create_operation(Constant, add_oper, 0, value=2)
    const_oper_2 = create_operation(Constant, add_oper, 1, value=3)

    return add_oper

@pytest.fixture
def large_operation_tree():
    add_oper = Addition()
    add_oper_2 = Addition()

    const_oper = create_operation(Constant, add_oper, 0, value=2)
    const_oper_2 = create_operation(Constant, add_oper, 1, value=3)

    const_oper_3 = create_operation(Constant, add_oper_2, 0, value=4)
    const_oper_4 = create_operation(Constant, add_oper_2, 1, value=5)

    add_oper_3 = Addition()
    add_oper_signal = Signal(add_oper, add_oper_3)
    add_oper._output_ports[0].connect(add_oper_signal)
    add_oper_3._input_ports[0].connect(add_oper_signal)

    add_oper_2_signal = Signal(add_oper_2, add_oper_3)
    add_oper_2._output_ports[0].connect(add_oper_2_signal)
    add_oper_3._input_ports[1].connect(add_oper_2_signal)
    return const_oper

def test_traverse_single_tree(operation):
    traverse = Traverse(operation)
    assert traverse.traverse() == [operation]

def test_traverse_tree(operation_tree):
    traverse = Traverse(operation_tree)
    assert len(traverse.traverse()) == 3

def test_traverse_large_tree(large_operation_tree):
    traverse = Traverse(large_operation_tree)
    assert len(traverse.traverse()) == 7

def test_traverse_type(large_operation_tree):
    traverse = Traverse(large_operation_tree)
    assert len(traverse.traverse(Addition)) == 3
    assert len(traverse.traverse(Constant)) == 4

def test_traverse_loop(operation_tree):
    add_oper_signal = Signal()
    operation_tree._output_ports[0].connect(add_oper_signal)
    operation_tree._input_ports[0].connect(add_oper_signal)
    traverse = Traverse(operation_tree)
    assert len(traverse.traverse()) == 2