From 68d8b8487dad1b87109042b2f80802a41937c090 Mon Sep 17 00:00:00 2001
From: Max Ng <maxgg@meta.com>
Date: Fri, 20 Sep 2024 16:45:36 -0700
Subject: [PATCH] Migrate verifyPfcWithMapChanges test cases

Summary:
- verifyPfcWithMapChanges_0
- verifyPfcWithMapChanges_1

Reviewed By: nivinl

Differential Revision: D62981045

fbshipit-source-id: 0700878ace597aec66ff949d0be56db8c21cbb12
---
 .../agent_hw_tests/AgentTrafficPfcTests.cpp   | 29 +++++++++++++++++--
 fboss/agent/test/utils/PfcTestUtils.cpp       | 16 +++++-----
 fboss/agent/test/utils/PfcTestUtils.h         |  1 +
 3 files changed, 36 insertions(+), 10 deletions(-)

diff --git a/fboss/agent/test/agent_hw_tests/AgentTrafficPfcTests.cpp b/fboss/agent/test/agent_hw_tests/AgentTrafficPfcTests.cpp
index 37630f235f48d..483ba39f3ca57 100644
--- a/fboss/agent/test/agent_hw_tests/AgentTrafficPfcTests.cpp
+++ b/fboss/agent/test/agent_hw_tests/AgentTrafficPfcTests.cpp
@@ -275,6 +275,7 @@ class AgentTrafficPfcTest : public AgentHwTest {
   void runTestWithCfg(
       const int trafficClass,
       const int pfcPriority,
+      const std::map<int, int>& tcToPgOverride = {},
       TrafficTestParams testParams = TrafficTestParams{},
       std::function<void(
           AgentEnsemble* ensemble,
@@ -301,7 +302,11 @@ class AgentTrafficPfcTest : public AgentHwTest {
         portIdsToConfigure = masterLogicalInterfacePortIds();
       }
       setupPfcBuffers(
-          cfg, portIdsToConfigure, kLosslessPgIds, testParams.buffer);
+          cfg,
+          portIdsToConfigure,
+          kLosslessPgIds,
+          tcToPgOverride,
+          testParams.buffer);
       applyNewConfig(cfg);
 
       setupEcmpTraffic(portIds);
@@ -374,7 +379,27 @@ TEST_P(AgentTrafficPfcGenTest, verifyPfc) {
   const int trafficClass = kLosslessTrafficClass;
   const int pfcPriority = kLosslessPriority;
   TrafficTestParams trafficParams = GetParam();
-  runTestWithCfg(trafficClass, pfcPriority, trafficParams);
+  runTestWithCfg(trafficClass, pfcPriority, {}, trafficParams);
+}
+
+// intent of this test is to send traffic so that it maps to
+// tc 2, now map tc 2 to PG 3. Mapping from PG to pfc priority
+// is 1:1, which means PG 3 is mapped to pfc priority 3.
+// Generate traffic to fire off PFC with smaller shared buffer
+TEST_F(AgentTrafficPfcTest, verifyPfcWithMapChanges_0) {
+  const int trafficClass = kLosslessTrafficClass;
+  const int pfcPriority = 3;
+  runTestWithCfg(trafficClass, pfcPriority, {{trafficClass, pfcPriority}});
+}
+
+// intent of this test is to send traffic so that it maps to
+// tc 7. Now we map tc 7 -> PG 2. Mapping from PG to pfc
+// priority is 1:1, which means PG 2 is mapped to pfc priority 2.
+// Generate traffic to fire off PFC with smaller shared buffer
+TEST_F(AgentTrafficPfcTest, verifyPfcWithMapChanges_1) {
+  const int trafficClass = 7;
+  const int pfcPriority = kLosslessPriority;
+  runTestWithCfg(trafficClass, pfcPriority, {{trafficClass, pfcPriority}});
 }
 
 } // namespace facebook::fboss
diff --git a/fboss/agent/test/utils/PfcTestUtils.cpp b/fboss/agent/test/utils/PfcTestUtils.cpp
index 49a4b92f6f109..9d0a977d63f9e 100644
--- a/fboss/agent/test/utils/PfcTestUtils.cpp
+++ b/fboss/agent/test/utils/PfcTestUtils.cpp
@@ -15,8 +15,7 @@ static const std::vector<int> kLossyPgIds{0};
 
 void setupQosMapForPfc(
     cfg::QosMap& qosMap,
-    const std::map<int, int>& tc2PgOverride = {},
-    const std::map<int, int>& pfcPri2PgIdOverride = {}) {
+    const std::map<int, int>& tc2PgOverride = {}) {
   // update pfc maps
   std::map<int16_t, int16_t> tc2PgId;
   std::map<int16_t, int16_t> tc2QueueId;
@@ -36,9 +35,6 @@ void setupQosMapForPfc(
   for (auto& tc2Pg : tc2PgOverride) {
     tc2PgId[tc2Pg.first] = tc2Pg.second;
   }
-  for (auto& tmp : pfcPri2PgIdOverride) {
-    pfcPri2PgId[tmp.first] = tmp.second;
-  }
 
   qosMap.dscpMaps()->resize(8);
   for (auto i = 0; i < 8; i++) {
@@ -53,7 +49,10 @@ void setupQosMapForPfc(
   qosMap.pfcPriorityToQueueId() = std::move(pfcPri2QueueId);
 }
 
-void setupPfc(cfg::SwitchConfig& cfg, const std::vector<PortID>& ports) {
+void setupPfc(
+    cfg::SwitchConfig& cfg,
+    const std::vector<PortID>& ports,
+    const std::map<int, int>& tcToPgOverride) {
   cfg::PortPfc pfc;
   pfc.tx() = true;
   pfc.rx() = true;
@@ -61,7 +60,7 @@ void setupPfc(cfg::SwitchConfig& cfg, const std::vector<PortID>& ports) {
 
   cfg::QosMap qosMap;
   // setup qos map with pfc structs
-  setupQosMapForPfc(qosMap);
+  setupQosMapForPfc(qosMap, tcToPgOverride);
 
   // setup qosPolicy
   cfg.qosPolicies()->resize(1);
@@ -152,8 +151,9 @@ void setupPfcBuffers(
     cfg::SwitchConfig& cfg,
     const std::vector<PortID>& ports,
     const std::vector<int>& losslessPgIds,
+    const std::map<int, int>& tcToPgOverride,
     PfcBufferParams buffer) {
-  setupPfc(cfg, ports);
+  setupPfc(cfg, ports, tcToPgOverride);
 
   std::map<std::string, std::vector<cfg::PortPgConfig>> portPgConfigMap;
   setupPortPgConfig(
diff --git a/fboss/agent/test/utils/PfcTestUtils.h b/fboss/agent/test/utils/PfcTestUtils.h
index 17a050061cecc..131f4661d7e41 100644
--- a/fboss/agent/test/utils/PfcTestUtils.h
+++ b/fboss/agent/test/utils/PfcTestUtils.h
@@ -21,6 +21,7 @@ void setupPfcBuffers(
     cfg::SwitchConfig& cfg,
     const std::vector<PortID>& ports,
     const std::vector<int>& losslessPgIds,
+    const std::map<int, int>& tcToPgOverride = {},
     PfcBufferParams buffer = PfcBufferParams{});
 
 void addPuntPfcPacketAcl(cfg::SwitchConfig& cfg, uint16_t queueId);