Add a fixture to enable nat for dpus

tests/common/helpers/dut_utils.py

@@ -21,6 +21,7 @@
 
 CONTAINER_CHECK_INTERVAL_SECS = 1
 CONTAINER_RESTART_THRESHOLD_SECS = 180
+NAT_ENABLE_KEY = "nat_enabled_on_{}"
 
 # Ansible config files
 LAB_CONNECTION_GRAPH_PATH = os.path.normpath((os.path.join(os.path.dirname(__file__), "../../../ansible/files")))
@@ -655,3 +656,66 @@ def extract_techsupport_tarball_file(duthost, tarball_name):
         techsupport_folder = tarball_name.split('.')[0].split('/var/dump/')[1]
         techsupport_folder_full_path = '{}{}'.format(dst_folder, techsupport_folder)
     return techsupport_folder_full_path
+
+
+def is_enabled_nat_for_dpu(duthost, request):
+    if request.config.cache.get(NAT_ENABLE_KEY.format(duthost.hostname), False):
+        logger.info("NAT is enabled")
+        return True
+    else:
+        logger.info('NAT is not enabled')
+        return False
+
+
+def get_dpu_names_and_ssh_ports(duthost, dpuhost_names, ansible_adhoc):
+    dpuhost_ssh_port_dict = {}
+    for dpuhost_name in dpuhost_names:
+        host = ansible_adhoc(become=True, args=[], kwargs={})[dpuhost_name]
+        vm = host.options["inventory_manager"].get_host(dpuhost_name).vars
+        ansible_ssh_port = vm.get("ansible_ssh_port", None)
+        if ansible_ssh_port:
+            dpuhost_ssh_port_dict[dpuhost_name] = ansible_ssh_port
+
+    duthost_name = duthost.hostname
+    dpu_name_ssh_port_dict = {}
+    for dpuhost_name, dpu_host_ssh_port in dpuhost_ssh_port_dict.items():
+        if duthost_name in dpuhost_name:
+            res = re.match(fr"{duthost_name}.*dpu.*(\d+)", dpuhost_name)
+            if res:
+                dpuhost_index = res[1]
+            else:
+                assert f"Not find the dpu name index in the {dpuhost_name}, please correct the dpuhost_name. " \
+                       f"dpuhost name should include the dut host name and the dpu index. " \
+                       f"e.g smartswitch-01-dpu-1, smartswitch-01 is the duthost name, " \
+                       f"dpu-1 is the dpu name, and 1 is the dpu index"
+            dpu_name_ssh_port_dict[f"dpu{dpuhost_index}"] = str(dpu_host_ssh_port)
+    logger.info(f"dpu_name_ssh_port_dict:{dpu_name_ssh_port_dict}")
+
+    return dpu_name_ssh_port_dict
+
+
+def check_nat_is_enabled_and_set_cache(duthost, request):
+    get_nat_iptable_output = 'sudo iptables -t nat -L'
+    nat_iptable_output = duthost.shell(get_nat_iptable_output)['stdout']
+    pattern_nat_result = '.*DNAT.*tcp.*anywhere.*anywhere.*tcp dpt:.* to:169.254.200.*22.*'
+    if re.search(pattern_nat_result, nat_iptable_output):
+        logger.info('NAT is enabled successfully')
+        request.config.cache.set(NAT_ENABLE_KEY.format(duthost.hostname), True)
+        return True
+    else:
+        raise Exception('NAT is not enabled successfully')
+
+
+def enable_nat_for_dpus(duthost, dpu_name_ssh_port_dict, request):
+    enable_nat_cmds = [
+        "sudo su",
+        "sudo sed -i 's/#net.ipv4.ip_forward=1/net.ipv4.ip_forward=1/g' /etc/sysctl.conf",
+        "sudo echo net.ipv4.conf.eth0.forwarding=1 >> /etc/sysctl.conf",
+        "sudo sysctl -p",
+        f"sudo sonic-dpu-mgmt-traffic.sh inbound -e --dpus "
+        f"{','.join(dpu_name_ssh_port_dict.keys())} --ports {','.join(dpu_name_ssh_port_dict.values())}",
+        "sudo iptables-save > /etc/iptables/rules.v4",
+        "exit"
+    ]
+    duthost.shell_cmds(cmds=enable_nat_cmds)
+    check_nat_is_enabled_and_set_cache(duthost, request)

tests/conftest.py

@@ -50,7 +50,8 @@
 from tests.common.utilities import str2bool
 from tests.common.utilities import safe_filename
 from tests.common.utilities import get_duts_from_host_pattern
-from tests.common.helpers.dut_utils import is_supervisor_node, is_frontend_node, create_duthost_console, creds_on_dut
+from tests.common.helpers.dut_utils import is_supervisor_node, is_frontend_node, create_duthost_console, creds_on_dut, \
+    is_enabled_nat_for_dpu, get_dpu_names_and_ssh_ports, enable_nat_for_dpus
 from tests.common.cache import FactsCache
 from tests.common.config_reload import config_reload
 from tests.common.helpers.assertions import pytest_assert as pt_assert
@@ -484,8 +485,27 @@ def duthost(duthosts, request):
     return duthost
 
 
+@pytest.fixture(scope="session")
+def enable_nat_for_dpuhosts(duthosts, ansible_adhoc, request):
+    """
+    @summary: fixture to enable nat for dpuhost.
+    @param duthosts: fixture to get DUT hosts
+    @param ansible_adhoc: Fixture provided by the pytest-ansible package.
+        Source of the various device objects. It is
+        mandatory argument for the class constructors.
+    @param request: request parameters for duthost test fixture
+    """
+    dpuhost_names = get_specified_dpus(request)
+    if dpuhost_names:
+        logging.info(f"dpuhost_names: {dpuhost_names}")
+        for duthost in duthosts:
+            if not is_enabled_nat_for_dpu(duthost, request):
+                dpu_name_ssh_port_dict = get_dpu_names_and_ssh_ports(duthost, dpuhost_names, ansible_adhoc)
+                enable_nat_for_dpus(duthost, dpu_name_ssh_port_dict, request)
+
+
 @pytest.fixture(name="dpuhosts", scope="session")
-def fixture_dpuhosts(enhance_inventory, ansible_adhoc, tbinfo, request):
+def fixture_dpuhosts(enhance_inventory, ansible_adhoc, tbinfo, request, enable_nat_for_dpuhosts):
     """
     @summary: fixture to get DPU hosts defined in testbed.
     @param ansible_adhoc: Fixture provided by the pytest-ansible package.