From 2cc91bf7346a99ad8171cc097b91d7ea6feaf4cc Mon Sep 17 00:00:00 2001
From: Stefano Avallone <stavallo@gmail.com>
Date: Fri, 26 Jul 2024 17:43:26 +0200
Subject: [PATCH] wifi: Extend MLO test to check Txop Link entity swapping

---
 src/wifi/test/wifi-mlo-test.cc | 63 ++++++++++++++++++++++++++++------
 src/wifi/test/wifi-mlo-test.h  | 17 +++++++++
 2 files changed, 70 insertions(+), 10 deletions(-)

diff --git a/src/wifi/test/wifi-mlo-test.cc b/src/wifi/test/wifi-mlo-test.cc
index a7e92c059..6e355f22f 100644
--- a/src/wifi/test/wifi-mlo-test.cc
+++ b/src/wifi/test/wifi-mlo-test.cc
@@ -150,6 +150,16 @@ GetRnrLinkInfoTest::DoRun()
                           "Unexpected tbtt ID of the second reported AP");
 }
 
+TypeId
+MldSwapLinksTest::TestWifiMac::GetTypeId()
+{
+    static TypeId tid = TypeId("ns3::TestWifiMac")
+                            .SetParent<WifiMac>()
+                            .SetGroupName("Wifi")
+                            .AddConstructor<TestWifiMac>();
+    return tid;
+}
+
 MldSwapLinksTest::MldSwapLinksTest()
     : TestCase("Test the WifiMac::SwapLinks() method")
 {
@@ -161,33 +171,62 @@ MldSwapLinksTest::RunOne(std::string text,
                          const std::map<uint8_t, uint8_t>& links,
                          const std::map<uint8_t, uint8_t>& expected)
 {
-    TestWifiMac mac;
+    auto mac = CreateObjectWithAttributes<TestWifiMac>("QosSupported",
+                                                       BooleanValue(false),
+                                                       "Txop",
+                                                       PointerValue(CreateObject<Txop>()));
 
     std::vector<Ptr<WifiPhy>> phys;
-    for (std::size_t i = 0; i < nLinks; i++)
+    std::vector<Ptr<FrameExchangeManager>> feManagers;
+
+    for (std::size_t i = 0; i < nLinks; ++i)
     {
-        phys.emplace_back(CreateObject<SpectrumWifiPhy>());
+        auto phy = CreateObject<SpectrumWifiPhy>();
+        phy->SetPhyId(i);
+        phys.emplace_back(phy);
+        feManagers.emplace_back(CreateObject<TestFrameExchangeManager>());
     }
-    mac.SetWifiPhys(phys); // create links containing the given PHYs
+    mac->SetWifiPhys(phys); // create links containing the given PHYs
+    mac->SetFrameExchangeManagers(feManagers);
+    mac->GetTxop()->SetWifiMac(mac);
 
-    mac.SwapLinks(links);
+    // set CWmin of each Txop LinkEntity to the link ID, so that we can check where it has moved
+    for (std::size_t id = 0; id < nLinks; ++id)
+    {
+        mac->GetTxop()->SetMinCw(id, id);
+    }
 
-    NS_TEST_EXPECT_MSG_EQ(mac.GetNLinks(), nLinks, "Number of links changed after swapping");
+    mac->SwapLinks(links);
+
+    NS_TEST_EXPECT_MSG_EQ(mac->GetNLinks(), nLinks, "Number of links changed after swapping");
 
     for (const auto& [linkId, phyId] : expected)
     {
-        NS_TEST_ASSERT_MSG_EQ(mac.GetLinks().count(linkId),
-                              1,
+        NS_TEST_ASSERT_MSG_EQ(mac->GetLinks().contains(linkId),
+                              true,
                               "Link ID " << +linkId << " does not exist");
 
         NS_TEST_ASSERT_MSG_LT(+phyId, nLinks, "Invalid PHY ID");
 
         // the id of the PHY operating on a link is the original ID of the link
-        NS_TEST_EXPECT_MSG_EQ(mac.GetWifiPhy(linkId),
-                              phys.at(phyId),
+        NS_TEST_EXPECT_MSG_EQ(+mac->GetWifiPhy(linkId)->GetPhyId(),
+                              +phyId,
                               text << ": Link " << +phyId << " has not been moved to link "
                                    << +linkId);
+
+        NS_TEST_EXPECT_MSG_EQ(
+            +DynamicCast<TestFrameExchangeManager>(mac->GetFrameExchangeManager(linkId))
+                 ->GetLinkId(),
+            +linkId,
+            text << ": Link ID stored by FrameExchangeManager has not been updated");
+
+        NS_TEST_EXPECT_MSG_EQ(mac->GetTxop()->GetMinCw(linkId),
+                              +phyId,
+                              text << ": Txop Link entity " << +phyId
+                                   << " has not been moved to link " << +linkId);
     }
+
+    mac->Dispose();
 }
 
 void
@@ -200,6 +239,10 @@ MldSwapLinksTest::DoRun()
            3,
            {{0, 2}, {2, 1}},
            {{0, 1}, {1, 2}, {2, 0}});
+    RunOne("A different non-circular swapping, same result",
+           3,
+           {{1, 0}, {2, 1}},
+           {{0, 1}, {1, 2}, {2, 0}});
     RunOne("One move only, autodetect how to complete the swapping",
            3,
            {{2, 0}},
diff --git a/src/wifi/test/wifi-mlo-test.h b/src/wifi/test/wifi-mlo-test.h
index 48f01fe24..513695449 100644
--- a/src/wifi/test/wifi-mlo-test.h
+++ b/src/wifi/test/wifi-mlo-test.h
@@ -21,6 +21,7 @@
 #define WIFI_MLO_TEST_H
 
 #include "ns3/ap-wifi-mac.h"
+#include "ns3/frame-exchange-manager.h"
 #include "ns3/mgt-action-headers.h"
 #include "ns3/mgt-headers.h"
 #include "ns3/multi-model-spectrum-channel.h"
@@ -74,6 +75,9 @@ class MldSwapLinksTest : public TestCase
       public:
         ~TestWifiMac() override = default;
 
+        /// \return the object TypeId
+        static TypeId GetTypeId();
+
         using WifiMac::GetLinks;
         using WifiMac::SwapLinks;
 
@@ -92,6 +96,19 @@ class MldSwapLinksTest : public TestCase
         }
     };
 
+    /**
+     * Test FrameExchangeManager subclass to access m_linkId
+     */
+    class TestFrameExchangeManager : public FrameExchangeManager
+    {
+      public:
+        /// \return the link ID stored by this object
+        uint8_t GetLinkId() const
+        {
+            return m_linkId;
+        }
+    };
+
   public:
     MldSwapLinksTest();
     ~MldSwapLinksTest() override = default;