from viridis import tree
from six.moves import range

import pytest

@pytest.fixture
def base_tree():
    t = tree.Ultrametric(list(range(6)))
    t.merge(0, 1, 0.1)
    t.merge(6, 2, 0.2)
    t.merge(3, 4, 0.3)
    t.merge(8, 5, 0.4)
    t.merge(7, 8, 0.5)
    return t


def test_split(base_tree):
    t = base_tree
    t.split(0, 2)
    assert t.node[9]['num_leaves'] == 3
    t.split(0, 4) # nothing to do
    assert t.node[9]['num_leaves'] == 3

from viridis import tree
from six.moves import range

def test_split():
    t = tree.Ultrametric(list(range(6)))
    t.merge(0, 1, 0.1)
    t.merge(6, 2, 0.2)
    t.merge(3, 4, 0.3)
    t.merge(8, 5, 0.4)
    t.merge(7, 8, 0.5)
    t.split(0, 2)
    assert t.node[9]['num_leaves'] == 3
    t.split(0, 4) # nothing to do
    assert t.node[9]['num_leaves'] == 3

--- 
+++ 
@@ -1,13 +1,21 @@
 from viridis import tree
 from six.moves import range
 
-def test_split():
+import pytest
+
+@pytest.fixture
+def base_tree():
     t = tree.Ultrametric(list(range(6)))
     t.merge(0, 1, 0.1)
     t.merge(6, 2, 0.2)
     t.merge(3, 4, 0.3)
     t.merge(8, 5, 0.4)
     t.merge(7, 8, 0.5)
+    return t
+
+
+def test_split(base_tree):
+    t = base_tree
     t.split(0, 2)
     assert t.node[9]['num_leaves'] == 3
     t.split(0, 4) # nothing to do
Update test to use pytest fixture
