Add tests for get_topological_weights

Author: pradyunsg

tests/unit/resolution_resolvelib/test_resolver.py

@@ -6,7 +6,10 @@
 
 from pip._internal.req.constructors import install_req_from_line
 from pip._internal.req.req_set import RequirementSet
-from pip._internal.resolution.resolvelib.resolver import Resolver
+from pip._internal.resolution.resolvelib.resolver import (
+    Resolver,
+    get_topological_weights,
+)
 
 
 @pytest.fixture()
@@ -87,3 +90,147 @@ def test_new_resolver_get_installation_order(resolver, edges, ordered_reqs):
     ireqs = resolver.get_installation_order(reqset)
     req_strs = [str(r.req) for r in ireqs]
     assert req_strs == ordered_reqs
+
+
+@pytest.mark.parametrize(
+    "name, edges, expected_weights",
+    [
+        (
+            # From https://github.com/pypa/pip/pull/8127#discussion_r414564664
+            "deep second edge",
+            [
+                (None, "one"),
+                (None, "two"),
+                ("one", "five"),
+                ("two", "three"),
+                ("three", "four"),
+                ("four", "five"),
+            ],
+            {None: 0, "one": 1, "two": 1, "three": 2, "four": 3, "five": 4},
+        ),
+        (
+            "linear",
+            [
+                (None, "one"),
+                ("one", "two"),
+                ("two", "three"),
+                ("three", "four"),
+                ("four", "five"),
+            ],
+            {None: 0, "one": 1, "two": 2, "three": 3, "four": 4, "five": 5},
+        ),
+        (
+            "linear AND root -> two",
+            [
+                (None, "one"),
+                ("one", "two"),
+                ("two", "three"),
+                ("three", "four"),
+                ("four", "five"),
+                (None, "two"),
+            ],
+            {None: 0, "one": 1, "two": 2, "three": 3, "four": 4, "five": 5},
+        ),
+        (
+            "linear AND root -> three",
+            [
+                (None, "one"),
+                ("one", "two"),
+                ("two", "three"),
+                ("three", "four"),
+                ("four", "five"),
+                (None, "three"),
+            ],
+            {None: 0, "one": 1, "two": 2, "three": 3, "four": 4, "five": 5},
+        ),
+        (
+            "linear AND root -> four",
+            [
+                (None, "one"),
+                ("one", "two"),
+                ("two", "three"),
+                ("three", "four"),
+                ("four", "five"),
+                (None, "four"),
+            ],
+            {None: 0, "one": 1, "two": 2, "three": 3, "four": 4, "five": 5},
+        ),
+        (
+            "linear AND root -> five",
+            [
+                (None, "one"),
+                ("one", "two"),
+                ("two", "three"),
+                ("three", "four"),
+                ("four", "five"),
+                (None, "five"),
+            ],
+            {None: 0, "one": 1, "two": 2, "three": 3, "four": 4, "five": 5},
+        ),
+        (
+            "linear AND one -> four",
+            [
+                (None, "one"),
+                ("one", "two"),
+                ("two", "three"),
+                ("three", "four"),
+                ("four", "five"),
+                ("one", "four"),
+            ],
+            {None: 0, "one": 1, "two": 2, "three": 3, "four": 4, "five": 5},
+        ),
+        (
+            "linear AND two -> four",
+            [
+                (None, "one"),
+                ("one", "two"),
+                ("two", "three"),
+                ("three", "four"),
+                ("four", "five"),
+                ("two", "four"),
+            ],
+            {None: 0, "one": 1, "two": 2, "three": 3, "four": 4, "five": 5},
+        ),
+        (
+            "linear AND four -> one (cycle)",
+            [
+                (None, "one"),
+                ("one", "two"),
+                ("two", "three"),
+                ("three", "four"),
+                ("four", "five"),
+                ("four", "one"),
+            ],
+            {None: 0, "one": 1, "two": 2, "three": 3, "four": 4, "five": 5},
+        ),
+        (
+            "linear AND four -> two (cycle)",
+            [
+                (None, "one"),
+                ("one", "two"),
+                ("two", "three"),
+                ("three", "four"),
+                ("four", "five"),
+                ("four", "two"),
+            ],
+            {None: 0, "one": 1, "two": 2, "three": 3, "four": 4, "five": 5},
+        ),
+        (
+            "linear AND four -> three (cycle)",
+            [
+                (None, "one"),
+                ("one", "two"),
+                ("two", "three"),
+                ("three", "four"),
+                ("four", "five"),
+                ("four", "three"),
+            ],
+            {None: 0, "one": 1, "two": 2, "three": 3, "four": 4, "five": 5},
+        ),
+    ],
+)
+def test_new_resolver_topological_weights(name, edges, expected_weights):
+    graph = _make_graph(edges)
+
+    weights = get_topological_weights(graph)
+    assert weights == expected_weights