1 files changed, 20 insertions, 2 deletions

diff --git a/test/sca/test_tree.py b/test/sca/test_tree.py
index 3dbb960..5ee6324 100644
--- a/test/sca/test_tree.py
+++ b/test/sca/test_tree.py
@@ -1,3 +1,6 @@
+import io
+from contextlib import redirect_stdout
+
 from pyecsca.sca.re.tree import Tree, Map
 import pandas as pd
 
@@ -12,5 +15,20 @@ def test_build_tree():
     mapping2 = pd.DataFrame([(1, 0, 0), (2, 0, 0), (3, 0, 0)], columns=inputs2, index=cfgs)
     dmap2 = Map(None, mapping2)
     tree = Tree.build(set(cfgs), dmap1, dmap2)
-    print()
-    tree.render()
\ No newline at end of file
+    with redirect_stdout(io.StringIO()):
+        tree.render()
+        tree.describe()
+
+
+def test_expand_tree():
+    cfgs = ["a", "b", "c"]
+    inputs1 = [1, 2]
+    mapping1 = pd.DataFrame([(0, 4), (0, 3), (0, 4)], columns=inputs1, index=cfgs)
+    dmap1 = Map(None, mapping1)
+
+    inputs2 = ["f", "e", "d"]
+    mapping2 = pd.DataFrame([(1, 0, 0), (2, 0, 0), (3, 0, 0)], columns=inputs2, index=cfgs)
+    dmap2 = Map(None, mapping2)
+    tree = Tree.build(set(cfgs), dmap1)
+    res = tree.expand(dmap2)
+    assert res.height > tree.height