test_operation.py

import pytest

from b_asic import Constant, Addition, MAD, Butterfly

class TestTraverse:
    def test_traverse_single_tree(self, operation):
        """Traverse a tree consisting of one operation."""
        constant = Constant(None)
        assert list(constant.traverse()) == [constant]

    def test_traverse_tree(self, operation_tree):
        """Traverse a basic addition tree with two constants."""
        assert len(list(operation_tree.traverse())) == 5

    def test_traverse_large_tree(self, large_operation_tree):
        """Traverse a larger tree."""
        assert len(list(large_operation_tree.traverse())) == 13

    def test_traverse_type(self, large_operation_tree):
        result = list(large_operation_tree.traverse())
        assert len(list(filter(lambda type_: isinstance(type_, Addition), result))) == 3
        assert len(list(filter(lambda type_: isinstance(type_, Constant), result))) == 4

    def test_traverse_loop(self, operation_graph_with_cycle):
        assert len(list(operation_graph_with_cycle.traverse())) == 8

class TestSplit:
    def test_split_mad(self):
        mad1 = MAD()
        split = mad1.split()
        assert len(split) == 1
        assert len(list(split[0].traverse())) == 10

    def test_split_butterfly(self):
        but1 = Butterfly()
        split = but1.split()
        assert len(split) == 2

class TestToSfg:
    def test_convert_mad_to_sfg(self):
        mad1 = MAD()
        mad1_sfg = mad1.to_sfg()

        assert mad1.evaluate(1,1,1) == mad1_sfg.evaluate(1,1,1)
        assert len(mad1_sfg.operations) == 6

    def test_butterfly_to_sfg(self):
        but1 = Butterfly()
        but1_sfg = but1.to_sfg()
        
        assert but1.evaluate(1,1) == but1_sfg.evaluate(1,1)
        assert len(but1_sfg.operations) == 6