diff --git a/tests/test_iqtree/test_build_tree.py b/tests/test_iqtree/test_build_tree.py
index 3983227c..9ff57102 100644
--- a/tests/test_iqtree/test_build_tree.py
+++ b/tests/test_iqtree/test_build_tree.py
@@ -30,11 +30,10 @@ def check_build_tree_model(
     expected = make_tree("(Human,Chimpanzee,(SpermWhale,HumpbackW));")
 
     got = piqtree.build_tree(aln, str(model) if coerce_str else model)
-
     # Check topology
     assert expected.same_topology(got.unrooted())
     # Check if branch lengths exist
-    assert all("length" in v.params for v in got.get_edge_vector())
+    assert all(v.length is not None for v in got.get_edge_vector(include_root=False))
 
 
 def check_build_tree(
diff --git a/tests/test_iqtree/test_protein_build_tree.py b/tests/test_iqtree/test_protein_build_tree.py
index d74c509a..17b009ca 100644
--- a/tests/test_iqtree/test_protein_build_tree.py
+++ b/tests/test_iqtree/test_protein_build_tree.py
@@ -18,9 +18,8 @@ def check_build_tree_model(
     coerce_str: bool = False,
 ) -> None:
     got = piqtree.build_tree(aln, str(model) if coerce_str else model, rand_seed=1)
-
     # Check if all branch lengths exist
-    assert all("length" in v.params for v in got.get_edge_vector())
+    assert all(v.length is not None for v in got.get_edge_vector(include_root=False))
 
 
 def check_build_tree(