Add ignore_same_node option to assert_no_duplicate_router_ids

pybatfish/client/asserts.py

@@ -114,23 +114,38 @@ def _subdict(d, keys):
     return {k: d.get(k) for k in keys}
 
 
-def _get_duplicate_router_ids(question_name, session=None, snapshot=None):
-    # type: (str, Optional[Session], Optional[str]) -> DataFrame
+def _get_duplicate_router_ids(
+    question_name: str,
+    session: Optional["Session"] = None,
+    snapshot: Optional[str] = None,
+    ignore_same_node: bool = False,
+) -> DataFrame:
     """Helper function to get rows with duplicate router IDs for a given protocol.
 
     :param question_name: The question name to be used to fetch the protocol process configuration
     :param session: Batfish session to use for asking the question
     :param snapshot: Snapshot on which to ask the question
+    :param ignore_same_node: whether to ignore duplicate router-ids on the same node
     """
     df = (
         getattr(_get_question_object(session, question_name), question_name)()
         .answer(snapshot)
         .frame()
     )
-    df_duplicate = df[df.duplicated(["Router_ID"], keep=False)].sort_values(
-        ["Router_ID"]
-    )
-
+    if ignore_same_node:
+        # Maps Router_ID to whether multiple nodes have that Router_ID
+        router_id_on_duplicate_nodes = (
+            df.drop_duplicates(["Node", "Router_ID"])
+            .value_counts(["Router_ID"])
+            .map(lambda x: x > 1)
+        )
+        df_duplicate = df[
+            df.apply(lambda x: router_id_on_duplicate_nodes[x["Router_ID"]][0], axis=1)
+        ].sort_values(["Router_ID"])
+    else:
+        df_duplicate = df[df.duplicated(["Router_ID"], keep=False)].sort_values(
+            ["Router_ID"]
+        )
     return df_duplicate
 
 
@@ -743,14 +758,14 @@ def assert_no_undefined_references(
 
 
 def assert_no_duplicate_router_ids(
-    snapshot=None,
-    nodes=None,
-    protocols=None,
-    soft=False,
-    session=None,
-    df_format="table",
-):
-    # type: (Optional[str], Optional[str], Optional[List[str]], bool, Optional[Session], str) -> bool
+    snapshot: Optional[str] = None,
+    nodes: Optional[str] = None,
+    protocols: Optional[List[str]] = None,
+    soft: bool = False,
+    session: Optional["Session"] = None,
+    df_format: str = "table",
+    ignore_same_node: bool = False,
+) -> bool:
     """Assert that there are no duplicate router IDs present in the snapshot.
 
     :param snapshot: the snapshot on which to check the assertion
@@ -760,6 +775,7 @@ def assert_no_duplicate_router_ids(
     :param session: Batfish session to use for the assertion
     :param df_format: How to format the Dataframe content in the output message.
         Valid options are 'table' and 'records' (each row is a key-value pairs).
+    :param ignore_same_node: whether to ignore duplicate router-ids on the same node
     """
     __tracebackhide__ = operator.methodcaller("errisinstance", BatfishAssertException)
 
@@ -782,7 +798,7 @@ def assert_no_duplicate_router_ids(
 
     for protocol in protocols_to_fetch:
         df_duplicate = _get_duplicate_router_ids(
-            protocol + "ProcessConfiguration", session, snapshot
+            protocol + "ProcessConfiguration", session, snapshot, ignore_same_node
         )
         if not df_duplicate.empty:
             found_duplicates = True

pybatfish/client/session.py

@@ -199,9 +199,14 @@ def assert_flows_succeed(
         )
 
     def assert_no_duplicate_router_ids(
-        self, snapshot=None, nodes=None, protocols=None, soft=False, df_format="table"
-    ):
-        # type: (Optional[str], Optional[str], Optional[List[str]], bool, str) -> bool
+        self,
+        snapshot: Optional[str] = None,
+        nodes: Optional[str] = None,
+        protocols: Optional[List[str]] = None,
+        soft: bool = False,
+        df_format: str = "table",
+        ignore_same_node: bool = False,
+    ) -> bool:
         """Assert that there are no duplicate router IDs present in the snapshot.
 
         :param snapshot: the snapshot on which to check the assertion
@@ -211,9 +216,16 @@ def assert_no_duplicate_router_ids(
             not a failure)
         :param df_format: How to format the Dataframe content in the output message.
             Valid options are 'table' and 'records' (each row is a key-value pairs).
+        :param ignore_same_node whether to ignore duplicate router-ids on the same node
         """
         return assert_no_duplicate_router_ids(
-            snapshot, nodes, protocols, soft, self.session, df_format
+            snapshot,
+            nodes,
+            protocols,
+            soft,
+            self.session,
+            df_format,
+            ignore_same_node,
         )
 
     def assert_no_forwarding_loops(self, snapshot=None, soft=False, df_format="table"):

tests/client/test_assert.py

@@ -918,6 +918,40 @@ def test_no_duplicate_router_ids_no_session():
         assert mock_df_duplicate.to_string() in str(excinfo.value)
 
 
+def test_no_duplicate_router_ids_no_session_ignore_same_node():
+    """Test assert_no_duplicate_router_ids with ignore_same_node=True."""
+    with patch.object(
+        bfq, "bgpProcessConfiguration", create=True
+    ) as bgpProcessConfiguration:
+        # Test success
+        mock_df_unique = DataFrame.from_records(
+            [
+                {"Router_ID": "1.1.1.1", "Node": "n1", "VRF": "vrf1"},
+                {"Router_ID": "1.1.1.1", "Node": "n1", "VRF": "vrf2"},
+                {"Router_ID": "1.1.1.2", "Node": "n2", "VRF": "vrf1"},
+            ]
+        )
+        bgpProcessConfiguration.return_value = MockQuestion(
+            MockTableAnswer(mock_df_unique)
+        )
+        assert_no_duplicate_router_ids(protocols=["bgp"], ignore_same_node=True)
+        # Test failure
+        mock_df_duplicate = DataFrame.from_records(
+            [
+                {"Router_ID": "1.1.1.1", "Node": "n1", "VRF": "vrf1"},
+                {"Router_ID": "1.1.1.1", "Node": "n1", "VRF": "vrf2"},
+                {"Router_ID": "1.1.1.1", "Node": "n2", "VRF": "vrf1"},
+            ]
+        )
+        bgpProcessConfiguration.return_value = MockQuestion(
+            MockTableAnswer(mock_df_duplicate)
+        )
+        with pytest.raises(BatfishAssertException) as excinfo:
+            assert_no_duplicate_router_ids(protocols=["bgp"], ignore_same_node=True)
+        # Ensure found answer is printed
+        assert mock_df_duplicate.to_string() in str(excinfo.value)
+
+
 def test_no_forwarding_loops():
     """Confirm no-forwarding-loops assert passes and fails as expected when specifying a session."""
     bf = Session(load_questions=False, use_deprecated_workmgr_v1=False)