1 files changed, 34 insertions, 0 deletions
diff --git a/networkx/algorithms/tests/test_dag.py b/networkx/algorithms/tests/test_dag.py
index 8e1d7cc8..c0fac32c 100644
--- a/networkx/algorithms/tests/test_dag.py
+++ b/networkx/algorithms/tests/test_dag.py
@@ -478,6 +478,40 @@ class TestDAG:
         assert sorting == test_nodes
 
 
+def test_topological_generations():
+    G = nx.DiGraph(
+        {
+            1: [2, 3],
+            2: [4, 5],
+            3: [7],
+            4: [],
+            5: [6, 7],
+            6: [],
+            7: [],
+        }
+    ).reverse()
+    # order within each generation is inconsequential
+    generations = [sorted(gen) for gen in nx.topological_generations(G)]
+    expected = [[4, 6, 7], [3, 5], [2], [1]]
+    assert generations == expected
+
+    MG = nx.MultiDiGraph(G.edges)
+    MG.add_edge(2, 1)
+    generations = [sorted(gen) for gen in nx.topological_generations(MG)]
+    assert generations == expected
+
+
+def test_topological_generations_empty():
+    G = nx.DiGraph()
+    assert list(nx.topological_generations(G)) == []
+
+
+def test_topological_generations_cycle():
+    G = nx.DiGraph([[2, 1], [3, 1], [1, 2]])
+    with pytest.raises(nx.NetworkXUnfeasible):
+        list(nx.topological_generations(G))
+
+
 def test_is_aperiodic_cycle():
     G = nx.DiGraph()
     nx.add_cycle(G, [1, 2, 3, 4])