graph_drawing/tests/test_graph.py

# vim: set fileencoding=utf-8 :
import unittest

from node import Node, Leaf
from graph import generate_graph


class TestGraph(unittest.TestCase):

    def setUp(self):
        self.l0, self.l1, self.multi = Leaf(0), Leaf(1), Leaf('test')

    def tearDown(self):
        pass

    def test_simple_unary(self):
        uminus = Node('-', self.l1)
        g = generate_graph(uminus, Node)
        expect = self.strip("""
        -
        │
        1
        """)
        assert g == expect

    def test_simple_binary(self):
        plus = Node('+', self.l0, self.l1)
        g = generate_graph(plus, Node)
        expect = self.strip("""
         +
        ╭┴╮
        0 1
        """)
        assert g == expect

    def test_multichar_unary(self):
        uminus = Node('-', self.multi)
        g = generate_graph(uminus, Node)
        expect = self.strip("""
         -
         │
        test
        """)
        print g
        print expect
        assert g == expect

    def test_multichar_binary(self):
        plus = Node('+', self.multi, self.l1)
        g = generate_graph(plus, Node)
        expect = self.strip("""
            +
        ╭───┴╮
        test 1
        """)
        assert g == expect

    def test_function(self):
        exp = Leaf('x')
        inf = Leaf('oo')
        minus_inf = Node('-', inf)
        integral = Node('int', exp, minus_inf, inf)
        g = generate_graph(integral, Node)
        expect = self.strip("""
         int
        ╭─┼──╮
        x -  oo
          │
          oo
        """)
        assert g == expect

    def strip(self, str):
        return str.replace('\n        ', '\n')[1:-1]