feat: add wait_for_bgp option to reboot

tests/acl/test_acl.py

@@ -1321,8 +1321,7 @@ def post_setup_hook(self, dut, localhost, populate_vlan_arp_entries, tbinfo, con
 
         """
         dut.command("config save -y")
-        up_bgp_neighbors = dut.get_bgp_neighbors_per_asic("established")
-        reboot(dut, localhost, safe_reboot=True, check_intf_up_ports=True)
+        reboot(dut, localhost, safe_reboot=True, check_intf_up_ports=True, wait_for_bgp=True)
         # We need some additional delay on e1031
         if dut.facts["platform"] == "x86_64-cel_e1031-r0":
             time.sleep(240)
@@ -1338,9 +1337,6 @@ def post_setup_hook(self, dut, localhost, populate_vlan_arp_entries, tbinfo, con
             assert result, "Not all transceivers are detected or interfaces are up in {} seconds".format(
                 MAX_WAIT_TIME_FOR_INTERFACES)
 
-        pytest_assert(
-            wait_until(300, 10, 0, dut.check_bgp_session_state_all_asics, up_bgp_neighbors, "established"),
-            "All BGP sessions are not up after reboot, no point in continuing the test")
         # Delay 10 seconds for route convergence
         time.sleep(10)
 

tests/common/reboot.py

@@ -226,7 +226,7 @@ def execute_reboot_helper():
 def reboot(duthost, localhost, reboot_type='cold', delay=10,
            timeout=0, wait=0, wait_for_ssh=True, wait_warmboot_finalizer=False, warmboot_finalizer_timeout=0,
            reboot_helper=None, reboot_kwargs=None, return_after_reconnect=False,
-           safe_reboot=False, check_intf_up_ports=False):
+           safe_reboot=False, check_intf_up_ports=False, wait_for_bgp=False):
     """
     reboots DUT
     :param duthost: DUT host object
@@ -237,11 +237,13 @@ def reboot(duthost, localhost, reboot_type='cold', delay=10,
     :param wait: time to wait for DUT to initialize
     :param wait_for_ssh: Wait for SSH startup
     :param return_after_reconnect: Return from function as soon as SSH reconnects
-    :param wait_warmboot_finalizer=True: Wait for WARMBOOT_FINALIZER done
+    :param wait_warmboot_finalizer: Wait for WARMBOOT_FINALIZER done
+    :param warmboot_finalizer_timeout: Timeout for waiting WARMBOOT_FINALIZER
     :param reboot_helper: helper function to execute the power toggling
     :param reboot_kwargs: arguments to pass to the reboot_helper
     :param safe_reboot: arguments to wait DUT ready after reboot
     :param check_intf_up_ports: arguments to check interface after reboot
+    :param wait_for_bgp: arguments to wait for BGP after reboot
     :return:
     """
     assert not (safe_reboot and return_after_reconnect)
@@ -351,6 +353,13 @@ def reboot(duthost, localhost, reboot_type='cold', delay=10,
         assert float(dut_uptime.strftime("%s")) > float(dut_datetime.strftime("%s")), "Device {} did not reboot". \
             format(hostname)
 
+    if wait_for_bgp:
+        bgp_neighbors = duthost.get_bgp_neighbors_per_asic(state="all")
+        pytest_assert(
+            wait_until(wait + 300, 10, 0, duthost.check_bgp_session_state_all_asics, bgp_neighbors),
+            "Not all bgp sessions are established after reboot",
+        )
+
 
 def positive_uptime(duthost, dut_datetime):
     dut_uptime = duthost.get_up_time()