tests: factorize code in bgp_bmp_vrf

Factorize code in bgp_bmp_vrf Signed-off-by: Louis Scalbert <louis.scalbert@6wind.com>
FRRouting · Oct 24, 2024 · bc8f184 · bc8f184
1 parent d04d1d0
commit bc8f184
Showing 1 changed file with 32 additions and 54 deletions.
diff --git a/tests/topotests/bgp_bmp_vrf/test_bgp_bmp_vrf.py b/tests/topotests/bgp_bmp_vrf/test_bgp_bmp_vrf.py
@@ -305,7 +305,7 @@ def configure_prefixes(tgen, node, asn, safi, prefixes, vrf=None, update=True):
         tgen.gears[node].vtysh_cmd("".join(cmd))
 
 
-def unicast_prefixes(policy):
+def _test_prefixes(policy, step=1):
     """
     Setup the BMP  monitor policy, Add and withdraw ipv4/v6 prefixes.
     Check if the previous actions are logged in the BMP server with the right
@@ -314,62 +314,40 @@ def unicast_prefixes(policy):
     tgen = get_topogen()
     set_bmp_policy(tgen, "r1", 65501, "bmp1", "unicast", policy, vrf="vrf1")
 
-    update_seq()
-
     prefixes = ["172.31.0.15/32", "2111::1111/128"]
-    # add prefixes
-    configure_prefixes(tgen, "r2", 65502, "unicast", prefixes)
 
-    logger.info("checking for updated prefixes")
+    for type in ("update", "withdraw"):
+        update_seq()
 
-    for ipver in [4, 6]:
-        if UPDATE_EXPECTED_JSON:
-            continue
-        ref_file = "{}/r1/show-bgp-ipv{}-update-step1.json".format(CWD, ipver)
-        expected = json.loads(open(ref_file).read())
-
-        test_func = partial(
-            topotest.router_json_cmp,
-            tgen.gears["r1"],
-            f"show bgp vrf vrf1 ipv{ipver} json",
-            expected,
+        # add prefixes
+        configure_prefixes(
+            tgen, "r2", 65502, "unicast", prefixes, update=(type == "update")
         )
-        _, res = topotest.run_and_expect(test_func, None, count=30, wait=1)
-        assertmsg = f"r1: BGP IPv{ipver} convergence failed"
-        assert res is None, assertmsg
-
-    # check
-    test_func = partial(check_for_prefixes, prefixes, "update", policy, 1)
-    success, _ = topotest.run_and_expect(test_func, None, count=30, wait=1)
-    assert success, "Checking the updated prefixes has been failed !."
-
-    update_seq()
 
-    # withdraw prefixes
-    configure_prefixes(tgen, "r2", 65502, "unicast", prefixes, update=False)
+        logger.info(f"checking for prefixes {type}")
 
-    logger.info("checking for withdrawn prefixes")
-
-    for ipver in [4, 6]:
-        if UPDATE_EXPECTED_JSON:
-            continue
-        ref_file = "{}/r1/show-bgp-ipv{}-withdraw-step1.json".format(CWD, ipver)
-        expected = json.loads(open(ref_file).read())
-
-        test_func = partial(
-            topotest.router_json_cmp,
-            tgen.gears["r1"],
-            f"show bgp vrf vrf1 ipv{ipver} json",
-            expected,
-        )
-        _, res = topotest.run_and_expect(test_func, None, count=30, wait=1)
-        assertmsg = f"r1: BGP IPv{ipver} convergence failed"
-        assert res is None, assertmsg
-
-    # check
-    test_func = partial(check_for_prefixes, prefixes, "withdraw", policy, 1)
-    success, _ = topotest.run_and_expect(test_func, None, count=30, wait=1)
-    assert success, "Checking the withdrawn prefixes has been failed !."
+        for ipver in [4, 6]:
+            if UPDATE_EXPECTED_JSON:
+                continue
+            ref_file = "{}/r1/show-bgp-ipv{}-{}-step{}.json".format(
+                CWD, ipver, type, step
+            )
+            expected = json.loads(open(ref_file).read())
+
+            test_func = partial(
+                topotest.router_json_cmp,
+                tgen.gears["r1"],
+                f"show bgp vrf vrf1 ipv{ipver} json",
+                expected,
+            )
+            _, res = topotest.run_and_expect(test_func, None, count=30, wait=1)
+            assertmsg = f"r1: BGP IPv{ipver} convergence failed"
+            assert res is None, assertmsg
+
+        # check
+        test_func = partial(check_for_prefixes, prefixes, type, policy, step)
+        success, _ = topotest.run_and_expect(test_func, None, count=30, wait=1)
+        assert success, "Checking the updated prefixes has been failed !."
 
 
 def test_bmp_server_logging():
@@ -409,11 +387,11 @@ def test_bmp_bgp_unicast():
     Add/withdraw bgp unicast prefixes and check the bmp logs.
     """
     logger.info("*** Unicast prefixes pre-policy logging ***")
-    unicast_prefixes(PRE_POLICY)
+    _test_prefixes(PRE_POLICY)
     logger.info("*** Unicast prefixes post-policy logging ***")
-    unicast_prefixes(POST_POLICY)
+    _test_prefixes(POST_POLICY)
     logger.info("*** Unicast prefixes loc-rib logging ***")
-    unicast_prefixes(LOC_RIB)
+    _test_prefixes(LOC_RIB)
 
 
 def test_peer_down():