Add test for matrix option in nonisomorphic trees generator

2 files changed, 10 insertions, 2 deletions

diff --git a/networkx/generators/nonisomorphic_trees.py b/networkx/generators/nonisomorphic_trees.py
index c7c04912..a563063e 100644
--- a/networkx/generators/nonisomorphic_trees.py
+++ b/networkx/generators/nonisomorphic_trees.py
@@ -76,8 +76,7 @@ def number_of_nonisomorphic_trees(order):
     ----------
 
     """
-    length = sum(1 for _ in nonisomorphic_trees(order))
-    return length
+    return sum(1 for _ in nonisomorphic_trees(order))
 
 
 def _next_rooted_tree(predecessor, p=None):
diff --git a/networkx/generators/tests/test_nonisomorphic_trees.py b/networkx/generators/tests/test_nonisomorphic_trees.py
index 06c59087..eecb1118 100644
--- a/networkx/generators/tests/test_nonisomorphic_trees.py
+++ b/networkx/generators/tests/test_nonisomorphic_trees.py
@@ -47,3 +47,12 @@ class TestGeneratorNonIsomorphicTrees():
         assert_equal(sorted(f(3)[0].edges()), [(0, 1), (0, 2)])
         assert_equal(sorted(f(4)[0].edges()), [(0, 1), (0, 3), (1, 2)])
         assert_equal(sorted(f(4)[1].edges()), [(0, 1), (0, 2), (0, 3)])
+
+    def test_nonisomorphic_trees_matrix(self):
+        trees_2 = [[[0, 1], [1, 0]]]
+        assert_equal(list(nx.nonisomorphic_trees(2, create="matrix")), trees_2)
+        trees_3 = [[[0, 1, 1], [1, 0, 0], [1, 0, 0]]]
+        assert_equal(list(nx.nonisomorphic_trees(3, create="matrix")), trees_3)
+        trees_4 = [[[0, 1, 0, 1], [1, 0, 1, 0], [0, 1, 0, 0], [1, 0, 0, 0]],
+                   [[0, 1, 1, 1], [1, 0, 0, 0], [1, 0, 0, 0], [1, 0, 0, 0]]]
+        assert_equal(list(nx.nonisomorphic_trees(4, create="matrix")), trees_4)