From bdb93d5502b22a8149da163314b8a33188703c52 Mon Sep 17 00:00:00 2001
From: Stefano Avallone <stavallo@unina.it>
Date: Fri, 5 May 2023 11:10:48 +0200
Subject: [PATCH] wifi: Add a listening VHT STA in wifi-mac-ofdma test

---
 src/wifi/test/wifi-mac-ofdma-test.cc | 37 ++++++++++++++++++++++------
 1 file changed, 30 insertions(+), 7 deletions(-)

diff --git a/src/wifi/test/wifi-mac-ofdma-test.cc b/src/wifi/test/wifi-mac-ofdma-test.cc
index 1148fd369..dbfe88cc4 100644
--- a/src/wifi/test/wifi-mac-ofdma-test.cc
+++ b/src/wifi/test/wifi-mac-ofdma-test.cc
@@ -153,11 +153,17 @@ TestMultiUserScheduler::SelectTxFormat()
                                  ? GetMaxSizeOfQosNullAmpdu(m_trigger)
                                  : 3500; // allows aggregation of 2 MPDUs in TB PPDUs
 
-        Time duration =
-            WifiPhy::CalculateTxDuration(ampduSize,
-                                         txVector,
-                                         m_apMac->GetWifiPhy()->GetPhyBand(),
-                                         m_apMac->GetStaList(SINGLE_LINK_OP_ID).begin()->first);
+        auto staList = m_apMac->GetStaList(SINGLE_LINK_OP_ID);
+        // ignore non-HE stations
+        for (auto it = staList.begin(); it != staList.end();)
+        {
+            it = m_apMac->GetHeSupported(it->second) ? std::next(it) : staList.erase(it);
+        }
+
+        Time duration = WifiPhy::CalculateTxDuration(ampduSize,
+                                                     txVector,
+                                                     m_apMac->GetWifiPhy()->GetPhyBand(),
+                                                     staList.begin()->first);
 
         uint16_t length;
         std::tie(length, duration) = HePhy::ConvertHeTbPpduDurationToLSigLength(
@@ -200,7 +206,12 @@ TestMultiUserScheduler::SelectTxFormat()
     {
         // try to send a DL MU PPDU
         m_psduMap.clear();
-        const std::map<uint16_t, Mac48Address>& staList = m_apMac->GetStaList(SINGLE_LINK_OP_ID);
+        auto staList = m_apMac->GetStaList(SINGLE_LINK_OP_ID);
+        // ignore non-HE stations
+        for (auto it = staList.cbegin(); it != staList.cend();)
+        {
+            it = m_apMac->GetHeSupported(it->second) ? std::next(it) : staList.erase(it);
+        }
         NS_ABORT_MSG_IF(staList.size() != 4, "There must be 4 associated stations");
 
         /* Initialize TX params */
@@ -292,7 +303,12 @@ TestMultiUserScheduler::ComputeWifiTxVector()
     m_txVector.SetTxPowerLevel(
         GetWifiRemoteStationManager(SINGLE_LINK_OP_ID)->GetDefaultTxPowerLevel());
 
-    const std::map<uint16_t, Mac48Address>& staList = m_apMac->GetStaList(SINGLE_LINK_OP_ID);
+    auto staList = m_apMac->GetStaList(SINGLE_LINK_OP_ID);
+    // ignore non-HE stations
+    for (auto it = staList.cbegin(); it != staList.cend();)
+    {
+        it = m_apMac->GetHeSupported(it->second) ? std::next(it) : staList.erase(it);
+    }
     NS_ABORT_MSG_IF(staList.size() != 4, "There must be 4 associated stations");
 
     HeRu::RuType ruType;
@@ -2053,6 +2069,13 @@ OfdmaAckSequenceTest::DoRun()
                                                          : WIFI_STANDARD_80211be);
     m_staDevices = NetDeviceContainer(m_staDevices, wifi.Install(phy, mac, wifiNewStaNodes));
 
+    // create a listening VHT station
+    wifi.SetStandard(WIFI_STANDARD_80211ac);
+    wifi.Install(phy, mac, Create<Node>());
+
+    wifi.SetStandard(m_scenario == WifiOfdmaScenario::HE ? WIFI_STANDARD_80211ax
+                                                         : WIFI_STANDARD_80211be);
+
     mac.SetType("ns3::ApWifiMac", "BeaconGeneration", BooleanValue(true));
     mac.SetMultiUserScheduler(
         "ns3::TestMultiUserScheduler",