From 09f7dc186cedac0a62769a401127276e84b3cc8a Mon Sep 17 00:00:00 2001
From: rraghav-cisco <58446052+rraghav-cisco@users.noreply.github.com>
Date: Sun, 22 Dec 2024 20:46:45 -0800
Subject: [PATCH] Adding a fixture to set scheduler to slower speeds and revert
 it back. (#15718)

Description of PR
Summary:
Fixes the flakiness of DWRR testcase. The PR adds a new fixture that slows down the scheduler without changing the underlying algorithm. This allows the dWRR test to pass consitently.

co-authorized by: jianquanye@microsoft.com
---
 tests/qos/qos_sai_base.py           | 57 +++++++++++++++++++++++++++++
 tests/qos/test_qos_sai.py           |  4 +-
 tests/saitests/py3/sai_qos_tests.py | 29 ++++++++++++++-
 3 files changed, 86 insertions(+), 4 deletions(-)

diff --git a/tests/qos/qos_sai_base.py b/tests/qos/qos_sai_base.py
index 62694a1bfea..ef6cb0ef6f5 100644
--- a/tests/qos/qos_sai_base.py
+++ b/tests/qos/qos_sai_base.py
@@ -2637,3 +2637,60 @@ def change_lag_lacp_timer(self, duthosts, get_src_dst_asic_and_duts, tbinfo, nbr
                     logger.info(
                         "Changing lacp timer multiplier to default for %s in %s" % (neighbor_lag_member, peer_device))
                     vm_host.no_lacp_time_multiplier(neighbor_lag_member)
+
+    def copy_and_run_set_cir_script_cisco_8000(self, dut, ports, asic="", speed="10000000"):
+        if dut.facts['asic_type'] != "cisco-8000":
+            raise RuntimeError("This function should have been called only for cisco-8000.")
+        dshell_script = '''
+from common import *
+from sai_utils import *
+
+def set_port_cir(interface, rate):
+    mp = get_mac_port(interface)
+    sch = mp.get_scheduler()
+    sch.set_credit_cir(rate)
+    sch.set_credit_eir_or_pir(rate, False)
+
+'''
+
+        for intf in ports:
+            dshell_script += f'\nset_port_cir("{intf}", {speed})'
+
+        script_path = "/tmp/set_scheduler.py"
+        dut.copy(content=dshell_script, dest=script_path)
+        dut.docker_copy_to_all_asics(
+            container_name=f"syncd{asic}",
+            src=script_path,
+            dst="/")
+
+    @pytest.fixture(scope="function", autouse=False)
+    def set_cir_change(self, get_src_dst_asic_and_duts, dutConfig):
+        dst_port = dutConfig['dutInterfaces'][dutConfig["testPorts"]["dst_port_id"]]
+        dst_dut = get_src_dst_asic_and_duts['dst_dut']
+        dst_asic = get_src_dst_asic_and_duts['dst_asic']
+        dst_index = dst_asic.asic_index
+
+        if dst_dut.facts['asic_type'] != "cisco-8000":
+            yield
+            return
+
+        interfaces = [dst_port]
+        output = dst_asic.shell(f"show interface portchannel | grep {dst_port}", module_ignore_errors=True)['stdout']
+        if output != '':
+            output = output.replace('(S)', '')
+            pattern = ' *[0-9]*  *PortChannel[0-9]*  *LACP\\(A\\)\\(Up\\)  *(Ethernet[0-9]*.*)'
+            import re
+            match = re.match(pattern, output)
+            if not match:
+                raise RuntimeError(f"Couldn't find required interfaces out of the output:{output}")
+            interfaces = match.group(1).split(' ')
+
+        # Set scheduler to 5 Gbps.
+        self.copy_and_run_set_cir_script_cisco_8000(
+            dut=dst_dut,
+            ports=interfaces,
+            asic=dst_index,
+            speed=5 * 1000 * 1000 * 1000)
+
+        yield
+        return
diff --git a/tests/qos/test_qos_sai.py b/tests/qos/test_qos_sai.py
index abd0cc23650..8983bace982 100644
--- a/tests/qos/test_qos_sai.py
+++ b/tests/qos/test_qos_sai.py
@@ -1540,7 +1540,7 @@ def testQosSaiDot1pPgMapping(
 
     def testQosSaiDwrr(
         self, ptfhost, duthosts, get_src_dst_asic_and_duts, dutTestParams, dutConfig, dutQosConfig, change_port_speed,
-        skip_src_dst_different_asic
+        skip_src_dst_different_asic, set_cir_change
     ):
         """
             Test QoS SAI DWRR
@@ -2093,7 +2093,7 @@ def testQosSaiSeparatedDscpToPgMapping(self, duthost, request, ptfhost,
 
     def testQosSaiDwrrWeightChange(
         self, get_src_dst_asic_and_duts, ptfhost, dutTestParams, dutConfig, dutQosConfig,
-        updateSchedProfile, skip_src_dst_different_asic
+        updateSchedProfile, skip_src_dst_different_asic, set_cir_change
     ):
         """
             Test QoS SAI DWRR runtime weight change
diff --git a/tests/saitests/py3/sai_qos_tests.py b/tests/saitests/py3/sai_qos_tests.py
index a2ffa360cac..bfd06eb47c5 100755
--- a/tests/saitests/py3/sai_qos_tests.py
+++ b/tests/saitests/py3/sai_qos_tests.py
@@ -3846,8 +3846,25 @@ def runTest(self):
                     break
                 recv_pkt = scapy.Ether(received.packet)
 
-        # Release port
-        self.sai_thrift_port_tx_enable(self.dst_client, asic_type, [dst_port_id], enable_port_by_unblock_queue=False)
+        if asic_type == 'cisco-8000':
+            cmd_opt = ""
+            if 'dst_asic_index' in self.test_params:
+                cmd_opt = "-n asic{}".format(self.test_params['dst_asic_index'])
+            cmd = "sudo show platform npu script {} -s set_scheduler.py".format(cmd_opt)
+            out, err, ret = self.exec_cmd_on_dut(
+                self.dst_server_ip,
+                self.test_params['dut_username'],
+                self.test_params['dut_password'],
+                cmd)
+            if err != "" and out == "":
+                raise RuntimeError("cmd({}) might have failed in the DUT. Error:{}".format(cmd, err))
+        else:
+            # Release port
+            self.sai_thrift_port_tx_enable(
+                self.dst_client,
+                asic_type,
+                [dst_port_id],
+                enable_port_by_unblock_queue=False)
 
         cnt = 0
         pkts = []
@@ -3872,6 +3889,14 @@ def runTest(self):
                 # Ignore captured non-IP packet
                 continue
 
+        if asic_type == 'cisco-8000':
+            # Release port
+            self.sai_thrift_port_tx_enable(
+                self.dst_client,
+                asic_type,
+                [dst_port_id],
+                enable_port_by_unblock_queue=False)
+
         queue_pkt_counters = [0] * (max(prio_list) + 1)
         queue_num_of_pkts = [0] * (max(prio_list) + 1)
         for prio, q_cnt in zip(prio_list, q_pkt_cnt):