From e1272139d6493a21931b631f5548d1912bc2fe5d Mon Sep 17 00:00:00 2001
From: Stefano Avallone <stavallo@gmail.com>
Date: Tue, 21 May 2024 16:12:52 +0200
Subject: [PATCH] wifi: Extend MLO test to check WifiMac::TidMappedOnLink()
 function

---
 src/wifi/test/wifi-mlo-test.cc | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/src/wifi/test/wifi-mlo-test.cc b/src/wifi/test/wifi-mlo-test.cc
index fd89b7ff6..8ac30e66a 100644
--- a/src/wifi/test/wifi-mlo-test.cc
+++ b/src/wifi/test/wifi-mlo-test.cc
@@ -1866,6 +1866,23 @@ MultiLinkSetupTest::CheckMlSetup()
                     "Incorrect link mapping stored by "
                         << (mac->GetTypeOfStation() == AP ? "AP" : "non-AP") << " MLD for " << dir
                         << " direction");
+
+                // check correctness of WifiMac::TidMappedOnLink function
+                std::set<uint8_t> setupLinks(m_setupLinks.cbegin(), m_setupLinks.cend());
+                for (uint8_t tid = 0; tid < 8; ++tid)
+                {
+                    const auto& linkSet = mapping.contains(tid) ? mapping.at(tid) : setupLinks;
+
+                    for (const auto linkId : setupLinks)
+                    {
+                        NS_TEST_EXPECT_MSG_EQ(
+                            mac->TidMappedOnLink(dest->GetAddress(), dir, tid, linkId),
+                            linkSet.contains(linkId),
+                            "Incorrect return value on " << (mac == m_apMac ? "AP" : "STA")
+                                                         << " direction " << dir << " TID " << +tid
+                                                         << " linkID " << +linkId);
+                    }
+                }
             }
         };