From b8808dcb9d4debdbd915766309155e3744184438 Mon Sep 17 00:00:00 2001
From: Stefano Avallone <stavallo@gmail.com>
Date: Sat, 10 May 2025 12:00:22 +0200
Subject: [PATCH] wifi: Add a unit test for WifiMacQueue flush operation

---
 .../model/wifi-mac-queue-scheduler-impl.h     |  4 +-
 src/wifi/test/wifi-mac-queue-test.cc          | 91 +++++++++++++++++++
 2 files changed, 94 insertions(+), 1 deletion(-)

diff --git a/src/wifi/model/wifi-mac-queue-scheduler-impl.h b/src/wifi/model/wifi-mac-queue-scheduler-impl.h
index 6309aff4d..a02fad65a 100644
--- a/src/wifi/model/wifi-mac-queue-scheduler-impl.h
+++ b/src/wifi/model/wifi-mac-queue-scheduler-impl.h
@@ -24,6 +24,7 @@
 #include <vector>
 
 class WifiMacQueueDropOldestTest;
+class WifiMacQueueFlushTest;
 
 namespace ns3
 {
@@ -42,8 +43,9 @@ template <class Priority, class Compare = std::less<Priority>>
 class WifiMacQueueSchedulerImpl : public WifiMacQueueScheduler
 {
   public:
-    /// allow WifiMacQueueDropOldestTest class access
+    /// allow test classes access
     friend class ::WifiMacQueueDropOldestTest;
+    friend class ::WifiMacQueueFlushTest;
 
     /**
      * @brief Get the type ID.
diff --git a/src/wifi/test/wifi-mac-queue-test.cc b/src/wifi/test/wifi-mac-queue-test.cc
index b25f2ea38..5f68ea461 100644
--- a/src/wifi/test/wifi-mac-queue-test.cc
+++ b/src/wifi/test/wifi-mac-queue-test.cc
@@ -417,6 +417,96 @@ WifiExtractExpiredMpdusTest::DoRun()
     Simulator::Destroy();
 }
 
+/**
+ * @ingroup wifi-test
+ * @ingroup tests
+ *
+ * @brief Test that a wifi MAC queue is correctly flushed even if (at least) a container queue
+ * is blocked.
+ */
+class WifiMacQueueFlushTest : public TestCase
+{
+  public:
+    WifiMacQueueFlushTest();
+
+  private:
+    void DoRun() override;
+};
+
+WifiMacQueueFlushTest::WifiMacQueueFlushTest()
+    : TestCase("Test WifiMacQueue Flush operation")
+{
+}
+
+void
+WifiMacQueueFlushTest::DoRun()
+{
+    auto wifiMacQueue = CreateObject<WifiMacQueue>(AC_BE);
+    auto wifiMacScheduler = CreateObject<FcfsWifiQueueScheduler>();
+    wifiMacScheduler->m_perAcInfo[AC_BE].wifiMacQueue = wifiMacQueue;
+    wifiMacQueue->SetScheduler(wifiMacScheduler);
+
+    const auto addr1 = Mac48Address::Allocate();
+    const auto addr2 = Mac48Address::Allocate();
+
+    // Initialize the queue with some packets.
+    const auto tid = wifiAcList.at(AC_BE).GetLowTid();
+    const std::size_t nPackets = 5;
+    for (std::size_t i = 0; i < nPackets; i++)
+    {
+        WifiMacHeader header(WIFI_MAC_QOSDATA);
+        header.SetAddr1(addr1);
+        header.SetAddr2(addr2);
+        header.SetQosTid(tid);
+        auto item = Create<WifiMpdu>(Create<Packet>(), header);
+        wifiMacQueue->Enqueue(item);
+    }
+
+    WifiContainerQueueId queueId(WifiContainerQueueType::WIFI_QOSDATA_QUEUE,
+                                 WifiRcvAddr::UNICAST,
+                                 addr1,
+                                 tid);
+
+    // Check that all elements are inserted successfully.
+    NS_TEST_EXPECT_MSG_EQ(wifiMacQueue->GetNPackets(),
+                          nPackets,
+                          "Queue has unexpected number of elements");
+
+    NS_TEST_EXPECT_MSG_EQ(wifiMacQueue->GetNPackets(queueId),
+                          nPackets,
+                          "Unexpected number of packets in container queue");
+
+    NS_TEST_EXPECT_MSG_EQ(wifiMacScheduler->GetNext(AC_BE, SINGLE_LINK_OP_ID).has_value(),
+                          true,
+                          "Expected the scheduler to return a container queue");
+
+    // block the queue containing the packets
+    wifiMacScheduler->BlockQueues(WifiQueueBlockedReason::TID_NOT_MAPPED,
+                                  AC_BE,
+                                  {WifiContainerQueueType::WIFI_QOSDATA_QUEUE},
+                                  addr1,
+                                  addr2,
+                                  {tid},
+                                  {SINGLE_LINK_OP_ID});
+
+    NS_TEST_EXPECT_MSG_EQ(wifiMacScheduler->GetNext(AC_BE, SINGLE_LINK_OP_ID).has_value(),
+                          false,
+                          "Expected the container queue to  be blocked");
+
+    wifiMacQueue->Flush();
+
+    // Check that all elements are removed successfully.
+    NS_TEST_EXPECT_MSG_EQ(wifiMacQueue->GetNPackets(),
+                          0,
+                          "Queue has unexpected number of elements");
+
+    NS_TEST_EXPECT_MSG_EQ(wifiMacQueue->GetNPackets(queueId),
+                          0,
+                          "Unexpected number of packets in container queue");
+
+    Simulator::Destroy();
+}
+
 /**
  * @ingroup wifi-test
  * @ingroup tests
@@ -434,6 +524,7 @@ WifiMacQueueTestSuite::WifiMacQueueTestSuite()
 {
     AddTestCase(new WifiMacQueueDropOldestTest, TestCase::Duration::QUICK);
     AddTestCase(new WifiExtractExpiredMpdusTest, TestCase::Duration::QUICK);
+    AddTestCase(new WifiMacQueueFlushTest, TestCase::Duration::QUICK);
 }
 
 static WifiMacQueueTestSuite g_wifiMacQueueTestSuite; ///< the test suite