Refactored TestTCP and TestUDP to simplify them. Pulled out the places where test needs access to private members and put those in a TestAccess class.

Author: feasel0

src/inet/TCPEndPoint.h

@@ -38,10 +38,6 @@
 
 namespace chip {
 
-namespace Transport {
-class TCPTest;
-};
-
 namespace Inet {
 
 class TCPTest;
@@ -529,7 +525,6 @@ class DLL_EXPORT TCPEndPoint : public EndPointBasis<TCPEndPoint>
     constexpr static size_t kMaxReceiveMessageSize = System::PacketBuffer::kMaxSizeWithoutReserve;
 
 protected:
-    friend class ::chip::Transport::TCPTest;
     friend class TCPTest;
 
     TCPEndPoint(EndPointManager<TCPEndPoint> & endPointManager) :

src/transport/raw/TCP.h

@@ -181,7 +181,7 @@ class DLL_EXPORT TCPBase : public Base
     void CloseActiveConnections();
 
 private:
-    friend class TCPTest;
+    friend class TCPTestAccess;
 
     /**
      * Find an active connection to the given peer or return nullptr if
@@ -283,7 +283,6 @@ class TCP : public TCPBase
     ~TCP() override { mPendingPackets.ReleaseAll(); }
 
 private:
-    friend class TCPTest;
     TCPBase::ActiveConnectionState mConnectionsBuffer[kActiveConnectionsSize];
     PoolImpl<PendingPacket, kPendingPacketSize, ObjectPoolMem::kInline, PendingPacketPoolType::Interface> mPendingPackets;
 };

src/transport/raw/tests/TestMessageHeader.cpp

@@ -23,14 +23,13 @@
  *
  */
 
+#include <gtest/gtest.h>
 #include <lib/core/ErrorStr.h>
 #include <lib/support/CHIPMem.h>
 #include <lib/support/CodeUtils.h>
 #include <protocols/Protocols.h>
 #include <transport/raw/MessageHeader.h>
 
-#include <gtest/gtest.h>
-
 namespace {
 
 using namespace chip;

src/transport/raw/tests/TestPeerAddress.cpp

@@ -21,13 +21,12 @@
 #include <stdint.h>
 #include <string.h>
 
+#include <gtest/gtest.h>
 #include <inet/IPAddress.h>
 #include <lib/core/DataModelTypes.h>
 #include <lib/core/PeerId.h>
 #include <transport/raw/PeerAddress.h>
 
-#include <gtest/gtest.h>
-
 namespace {
 
 using namespace chip;

src/transport/raw/tests/TestTCP.cpp

@@ -23,6 +23,7 @@
 
 #include "NetworkTestHelpers.h"
 
+#include <gtest/gtest.h>
 #include <lib/core/CHIPCore.h>
 #include <lib/core/CHIPEncoding.h>
 #include <lib/support/CHIPMem.h>
@@ -32,38 +33,55 @@
 #include <transport/TransportMgr.h>
 #include <transport/raw/TCP.h>
 
-#include <gtest/gtest.h>
-
 #include <errno.h>
 #include <stdlib.h>
 #include <string.h>
 #include <utility>
 
 using namespace chip;
+using namespace chip::Test;
 using namespace chip::Inet;
 
-namespace {
-
 constexpr size_t kMaxTcpActiveConnectionCount = 4;
 constexpr size_t kMaxTcpPendingPackets        = 4;
 constexpr uint16_t kPacketSizeBytes           = static_cast<uint16_t>(sizeof(uint16_t));
 
 using TCPImpl = Transport::TCP<kMaxTcpActiveConnectionCount, kMaxTcpPendingPackets>;
 
+namespace chip {
+namespace Transport {
+class TCPTestAccess
+{
+public:
+    static void * FindActiveConnection(TCPImpl & tcp, Transport::PeerAddress & lPeerAddress)
+    {
+        return tcp.FindActiveConnection(lPeerAddress);
+    }
+    static TCPEndPoint * GetEndpoint(void * state) { return static_cast<TCPBase::ActiveConnectionState *>(state)->mEndPoint; }
+
+    static CHIP_ERROR ProcessReceivedBuffer(TCPImpl & tcp, TCPEndPoint * endPoint, const PeerAddress & peerAddress,
+                                            System::PacketBufferHandle && buffer)
+    {
+        return tcp.ProcessReceivedBuffer(endPoint, peerAddress, std::move(buffer));
+    }
+};
+} // namespace Transport
+} // namespace chip
+
+namespace {
+
 constexpr NodeId kSourceNodeId      = 123654;
 constexpr NodeId kDestinationNodeId = 111222333;
 constexpr uint32_t kMessageCounter  = 18;
 
-using TestContext = chip::Test::IOContext;
-
 const char PAYLOAD[] = "Hello!";
 
 class MockTransportMgrDelegate : public chip::TransportMgrDelegate
 {
 public:
     typedef int (*MessageReceivedCallback)(const uint8_t * message, size_t length, int count, void * data);
 
-    MockTransportMgrDelegate(TestContext * inContext) : mContext(inContext), mCallback(nullptr), mCallbackData(nullptr) {}
+    MockTransportMgrDelegate(IOContext * inContext) : mContext(inContext), mCallback(nullptr), mCallbackData(nullptr) {}
     ~MockTransportMgrDelegate() override {}
 
     void SetCallback(MessageReceivedCallback callback = nullptr, void * callback_data = nullptr)
@@ -151,66 +169,12 @@ class MockTransportMgrDelegate : public chip::TransportMgrDelegate
     int mReceiveHandlerCallCount = 0;
 
 private:
-    TestContext * mContext;
+    IOContext * mContext;
     MessageReceivedCallback mCallback;
     void * mCallbackData;
     TransportMgrBase mTransportMgrBase;
 };
 
-/////////////////////////// Init test
-
-class TestTCP : public ::testing::Test, public chip::Test::IOContext
-{
-protected:
-    void SetUp() { ASSERT_EQ(Init(), CHIP_NO_ERROR); }
-    void TearDown() { Shutdown(); }
-
-    void CheckSimpleInitTest(Inet::IPAddressType type)
-    {
-        TCPImpl tcp;
-
-        CHIP_ERROR err = tcp.Init(Transport::TcpListenParameters(GetTCPEndPointManager()).SetAddressType(type));
-
-        EXPECT_EQ(err, CHIP_NO_ERROR);
-    }
-
-    void CheckMessageTest(const IPAddress & addr)
-    {
-        TCPImpl tcp;
-
-        MockTransportMgrDelegate gMockTransportMgrDelegate(this);
-        gMockTransportMgrDelegate.InitializeMessageTest(tcp, addr);
-        gMockTransportMgrDelegate.SingleMessageTest(tcp, addr);
-        gMockTransportMgrDelegate.FinalizeMessageTest(tcp, addr);
-    }
-};
-
-#if INET_CONFIG_ENABLE_IPV4
-TEST_F(TestTCP, CheckSimpleInitTest4)
-{
-    CheckSimpleInitTest(IPAddressType::kIPv4);
-}
-
-TEST_F(TestTCP, CheckMessageTest4)
-{
-    IPAddress addr;
-    IPAddress::FromString("127.0.0.1", addr);
-    CheckMessageTest(addr);
-}
-#endif
-
-TEST_F(TestTCP, CheckSimpleInitTest6)
-{
-    CheckSimpleInitTest(IPAddressType::kIPv6);
-}
-
-TEST_F(TestTCP, CheckMessageTest6)
-{
-    IPAddress addr;
-    IPAddress::FromString("::1", addr);
-    CheckMessageTest(addr);
-}
-
 // Generates a packet buffer or a chain of packet buffers for a single message.
 struct TestData
 {
@@ -338,108 +302,177 @@ void TestData::Free()
     mMessageOffset = 0;
 }
 
-int TestDataCallbackCheck(const uint8_t * message, size_t length, int count, void * data)
+// Test Context
+
+class TestTCP : public ::testing::Test
 {
-    if (data == nullptr)
-    {
-        return -1;
-    }
-    TestData * currentData = static_cast<TestData *>(data) + count;
-    if (currentData->mPayload == nullptr)
-    {
-        return -2;
-    }
-    if (currentData->mMessageLength != length)
+public:
+    // Performs shared setup for all tests in the test suite
+    static void SetUpTestSuite()
     {
-        return -3;
+        if (mIOContext == nullptr)
+        {
+            mIOContext = new IOContext();
+            ASSERT_NE(mIOContext, nullptr);
+        }
+        ASSERT_EQ(mIOContext->Init(), CHIP_NO_ERROR);
     }
-    if (memcmp(currentData->mPayload + currentData->mMessageOffset, message, length) != 0)
+
+    // Performs shared teardown for all tests in the test suite
+    static void TearDownTestSuite()
     {
-        return -4;
+        if (mIOContext != nullptr)
+        {
+            mIOContext->Shutdown();
+            delete mIOContext;
+            mIOContext = nullptr;
+        }
     }
-    return 0;
-}
 
-} // namespace
+protected:
+    static IOContext * mIOContext;
 
-namespace chip {
-namespace Transport {
-class TCPTest
-{
-public:
-    static void CheckProcessReceivedBuffer(TestContext * ctx)
+    void CheckSimpleInitTest(IPAddressType type)
     {
         TCPImpl tcp;
 
-        IPAddress addr;
-        IPAddress::FromString("::1", addr);
+        CHIP_ERROR err = tcp.Init(Transport::TcpListenParameters(mIOContext->GetTCPEndPointManager()).SetAddressType(type));
 
-        MockTransportMgrDelegate gMockTransportMgrDelegate(ctx);
-        gMockTransportMgrDelegate.InitializeMessageTest(tcp, addr);
+        EXPECT_EQ(err, CHIP_NO_ERROR);
+    }
+
+    void CheckMessageTest(const IPAddress & addr)
+    {
+        TCPImpl tcp;
 
-        // Send a packet to get TCP going, so that we can find a TCPEndPoint to pass to ProcessReceivedBuffer.
-        // (The current TCPEndPoint implementation is not effectively mockable.)
+        MockTransportMgrDelegate gMockTransportMgrDelegate(mIOContext);
+        gMockTransportMgrDelegate.InitializeMessageTest(tcp, addr);
         gMockTransportMgrDelegate.SingleMessageTest(tcp, addr);
+        gMockTransportMgrDelegate.FinalizeMessageTest(tcp, addr);
+    }
 
-        Transport::PeerAddress lPeerAddress    = Transport::PeerAddress::TCP(addr);
-        TCPBase::ActiveConnectionState * state = tcp.FindActiveConnection(lPeerAddress);
-        ASSERT_NE(state, nullptr);
-        Inet::TCPEndPoint * lEndPoint = state->mEndPoint;
-        ASSERT_NE(lEndPoint, nullptr);
+    static int TestDataCallbackCheck(const uint8_t * message, size_t length, int count, void * data)
+    {
+        if (data == nullptr)
+        {
+            return -1;
+        }
+        TestData * currentData = static_cast<TestData *>(data) + count;
+        if (currentData->mPayload == nullptr)
+        {
+            return -2;
+        }
+        if (currentData->mMessageLength != length)
+        {
+            return -3;
+        }
+        if (memcmp(currentData->mPayload + currentData->mMessageOffset, message, length) != 0)
+        {
+            return -4;
+        }
+        return 0;
+    }
+};
 
-        CHIP_ERROR err = CHIP_NO_ERROR;
-        TestData testData[2];
-        gMockTransportMgrDelegate.SetCallback(TestDataCallbackCheck, testData);
+IOContext * TestTCP::mIOContext = nullptr;
 
-        // Test a single packet buffer.
-        gMockTransportMgrDelegate.mReceiveHandlerCallCount = 0;
-        EXPECT_TRUE(testData[0].Init((const uint16_t[]){ 111, 0 }));
-        err = tcp.ProcessReceivedBuffer(lEndPoint, lPeerAddress, std::move(testData[0].mHandle));
-        EXPECT_EQ(err, CHIP_NO_ERROR);
-        EXPECT_EQ(gMockTransportMgrDelegate.mReceiveHandlerCallCount, 1);
+// IPv4 Tests
 
-        // Test a message in a chain of three packet buffers. The message length is split across buffers.
-        gMockTransportMgrDelegate.mReceiveHandlerCallCount = 0;
-        EXPECT_TRUE(testData[0].Init((const uint16_t[]){ 1, 122, 123, 0 }));
-        err = tcp.ProcessReceivedBuffer(lEndPoint, lPeerAddress, std::move(testData[0].mHandle));
-        EXPECT_EQ(err, CHIP_NO_ERROR);
-        EXPECT_EQ(gMockTransportMgrDelegate.mReceiveHandlerCallCount, 1);
-
-        // Test two messages in a chain.
-        gMockTransportMgrDelegate.mReceiveHandlerCallCount = 0;
-        EXPECT_TRUE(testData[0].Init((const uint16_t[]){ 131, 0 }));
-        EXPECT_TRUE(testData[1].Init((const uint16_t[]){ 132, 0 }));
-        testData[0].mHandle->AddToEnd(std::move(testData[1].mHandle));
-        err = tcp.ProcessReceivedBuffer(lEndPoint, lPeerAddress, std::move(testData[0].mHandle));
-        EXPECT_EQ(err, CHIP_NO_ERROR);
-        EXPECT_EQ(gMockTransportMgrDelegate.mReceiveHandlerCallCount, 2);
-
-        // Test a chain of two messages, each a chain.
-        gMockTransportMgrDelegate.mReceiveHandlerCallCount = 0;
-        EXPECT_TRUE(testData[0].Init((const uint16_t[]){ 141, 142, 0 }));
-        EXPECT_TRUE(testData[1].Init((const uint16_t[]){ 143, 144, 0 }));
-        testData[0].mHandle->AddToEnd(std::move(testData[1].mHandle));
-        err = tcp.ProcessReceivedBuffer(lEndPoint, lPeerAddress, std::move(testData[0].mHandle));
-        EXPECT_EQ(err, CHIP_NO_ERROR);
-        EXPECT_EQ(gMockTransportMgrDelegate.mReceiveHandlerCallCount, 2);
-
-        // Test a message that is too large to coalesce into a single packet buffer.
-        gMockTransportMgrDelegate.mReceiveHandlerCallCount = 0;
-        gMockTransportMgrDelegate.SetCallback(TestDataCallbackCheck, &testData[1]);
-        EXPECT_TRUE(testData[0].Init((const uint16_t[]){ 51, System::PacketBuffer::kMaxSizeWithoutReserve, 0 }));
-        // Sending only the first buffer of the long chain. This should be enough to trigger the error.
-        System::PacketBufferHandle head = testData[0].mHandle.PopHead();
-        err                             = tcp.ProcessReceivedBuffer(lEndPoint, lPeerAddress, std::move(head));
-        EXPECT_EQ(err, CHIP_ERROR_MESSAGE_TOO_LONG);
-        EXPECT_EQ(gMockTransportMgrDelegate.mReceiveHandlerCallCount, 0);
+#if INET_CONFIG_ENABLE_IPV4
+TEST_F(TestTCP, CheckSimpleInitTest4)
+{
+    CheckSimpleInitTest(IPAddressType::kIPv4);
+}
 
-        gMockTransportMgrDelegate.FinalizeMessageTest(tcp, addr);
-    }
-};
-} // namespace Transport
-} // namespace chip
+TEST_F(TestTCP, CheckMessageTest4)
+{
+    IPAddress addr;
+    IPAddress::FromString("127.0.0.1", addr);
+    CheckMessageTest(addr);
+}
+#endif
+
+// IPv6 Tests
+
+TEST_F(TestTCP, CheckSimpleInitTest6)
+{
+    CheckSimpleInitTest(IPAddressType::kIPv6);
+}
+
+TEST_F(TestTCP, CheckMessageTest6)
+{
+    IPAddress addr;
+    IPAddress::FromString("::1", addr);
+    CheckMessageTest(addr);
+}
 
 TEST_F(TestTCP, CheckProcessReceivedBuffer)
 {
-    chip::Transport::TCPTest::CheckProcessReceivedBuffer(this);
+    TCPImpl tcp;
+
+    IPAddress addr;
+    IPAddress::FromString("::1", addr);
+
+    MockTransportMgrDelegate gMockTransportMgrDelegate(mIOContext);
+    gMockTransportMgrDelegate.InitializeMessageTest(tcp, addr);
+
+    // Send a packet to get TCP going, so that we can find a TCPEndPoint to pass to ProcessReceivedBuffer.
+    // (The current TCPEndPoint implementation is not effectively mockable.)
+    gMockTransportMgrDelegate.SingleMessageTest(tcp, addr);
+
+    Transport::PeerAddress lPeerAddress = Transport::PeerAddress::TCP(addr);
+    void * state                        = Transport::TCPTestAccess::FindActiveConnection(tcp, lPeerAddress);
+    ASSERT_NE(state, nullptr);
+    TCPEndPoint * lEndPoint = Transport::TCPTestAccess::GetEndpoint(state);
+    ASSERT_NE(lEndPoint, nullptr);
+
+    CHIP_ERROR err = CHIP_NO_ERROR;
+    TestData testData[2];
+    gMockTransportMgrDelegate.SetCallback(TestDataCallbackCheck, testData);
+
+    // Test a single packet buffer.
+    gMockTransportMgrDelegate.mReceiveHandlerCallCount = 0;
+    EXPECT_TRUE(testData[0].Init((const uint16_t[]){ 111, 0 }));
+    err = Transport::TCPTestAccess::ProcessReceivedBuffer(tcp, lEndPoint, lPeerAddress, std::move(testData[0].mHandle));
+    EXPECT_EQ(err, CHIP_NO_ERROR);
+    EXPECT_EQ(gMockTransportMgrDelegate.mReceiveHandlerCallCount, 1);
+
+    // Test a message in a chain of three packet buffers. The message length is split across buffers.
+    gMockTransportMgrDelegate.mReceiveHandlerCallCount = 0;
+    EXPECT_TRUE(testData[0].Init((const uint16_t[]){ 1, 122, 123, 0 }));
+    err = Transport::TCPTestAccess::ProcessReceivedBuffer(tcp, lEndPoint, lPeerAddress, std::move(testData[0].mHandle));
+    EXPECT_EQ(err, CHIP_NO_ERROR);
+    EXPECT_EQ(gMockTransportMgrDelegate.mReceiveHandlerCallCount, 1);
+
+    // Test two messages in a chain.
+    gMockTransportMgrDelegate.mReceiveHandlerCallCount = 0;
+    EXPECT_TRUE(testData[0].Init((const uint16_t[]){ 131, 0 }));
+    EXPECT_TRUE(testData[1].Init((const uint16_t[]){ 132, 0 }));
+    testData[0].mHandle->AddToEnd(std::move(testData[1].mHandle));
+    err = Transport::TCPTestAccess::ProcessReceivedBuffer(tcp, lEndPoint, lPeerAddress, std::move(testData[0].mHandle));
+    EXPECT_EQ(err, CHIP_NO_ERROR);
+    EXPECT_EQ(gMockTransportMgrDelegate.mReceiveHandlerCallCount, 2);
+
+    // Test a chain of two messages, each a chain.
+    gMockTransportMgrDelegate.mReceiveHandlerCallCount = 0;
+    EXPECT_TRUE(testData[0].Init((const uint16_t[]){ 141, 142, 0 }));
+    EXPECT_TRUE(testData[1].Init((const uint16_t[]){ 143, 144, 0 }));
+    testData[0].mHandle->AddToEnd(std::move(testData[1].mHandle));
+    err = Transport::TCPTestAccess::ProcessReceivedBuffer(tcp, lEndPoint, lPeerAddress, std::move(testData[0].mHandle));
+    EXPECT_EQ(err, CHIP_NO_ERROR);
+    EXPECT_EQ(gMockTransportMgrDelegate.mReceiveHandlerCallCount, 2);
+
+    // Test a message that is too large to coalesce into a single packet buffer.
+    gMockTransportMgrDelegate.mReceiveHandlerCallCount = 0;
+    gMockTransportMgrDelegate.SetCallback(TestDataCallbackCheck, &testData[1]);
+    EXPECT_TRUE(testData[0].Init((const uint16_t[]){ 51, System::PacketBuffer::kMaxSizeWithoutReserve, 0 }));
+    // Sending only the first buffer of the long chain. This should be enough to trigger the error.
+    System::PacketBufferHandle head = testData[0].mHandle.PopHead();
+    err = Transport::TCPTestAccess::ProcessReceivedBuffer(tcp, lEndPoint, lPeerAddress, std::move(head));
+    EXPECT_EQ(err, CHIP_ERROR_MESSAGE_TOO_LONG);
+    EXPECT_EQ(gMockTransportMgrDelegate.mReceiveHandlerCallCount, 0);
+
+    gMockTransportMgrDelegate.FinalizeMessageTest(tcp, addr);
 }
+
+} // namespace

src/transport/raw/tests/TestUDP.cpp

@@ -23,16 +23,16 @@
 
 #include "NetworkTestHelpers.h"
 
+#include <gtest/gtest.h>
 #include <lib/core/CHIPCore.h>
 #include <lib/support/CodeUtils.h>
 #include <transport/TransportMgr.h>
 #include <transport/raw/UDP.h>
 
-#include <gtest/gtest.h>
-
 #include <errno.h>
 
 using namespace chip;
+using namespace chip::Test;
 using namespace chip::Inet;
 
 namespace {
@@ -41,8 +41,6 @@ constexpr NodeId kSourceNodeId      = 123654;
 constexpr NodeId kDestinationNodeId = 111222333;
 constexpr uint32_t kMessageCounter  = 18;
 
-using TestContext = chip::Test::IOContext;
-
 const char PAYLOAD[]        = "Hello!";
 int ReceiveHandlerCallCount = 0;
 
@@ -70,87 +68,114 @@ class MockTransportMgrDelegate : public TransportMgrDelegate
     }
 };
 
-} // namespace
-
-/////////////////////////// Init test
+// Test Context
 
 class TestUDP : public ::testing::Test
 {
+public:
+    // Performs shared setup for all tests in the test suite
+    static void SetUpTestSuite()
+    {
+        if (mIOContext == nullptr)
+        {
+            mIOContext = new IOContext();
+            ASSERT_NE(mIOContext, nullptr);
+        }
+        ASSERT_EQ(mIOContext->Init(), CHIP_NO_ERROR);
+    }
+
+    // Performs shared teardown for all tests in the test suite
+    static void TearDownTestSuite()
+    {
+        if (mIOContext != nullptr)
+        {
+            mIOContext->Shutdown();
+            delete mIOContext;
+            mIOContext = nullptr;
+        }
+    }
+
 protected:
-    TestUDP() { inContext = new TestContext(); }
-    ~TestUDP() { delete inContext; }
-    void SetUp() { ASSERT_EQ(inContext->Init(), CHIP_NO_ERROR); }
-    void TearDown() { inContext->Shutdown(); }
-    TestContext * inContext;
-};
+    static IOContext * mIOContext;
 
-void CheckSimpleInitTest(TestContext * ctx, Inet::IPAddressType type)
-{
-    Transport::UDP udp;
+    void CheckSimpleInitTest(IPAddressType type)
+    {
+        Transport::UDP udp;
 
-    CHIP_ERROR err = udp.Init(Transport::UdpListenParameters(ctx->GetUDPEndPointManager()).SetAddressType(type).SetListenPort(0));
+        CHIP_ERROR err =
+            udp.Init(Transport::UdpListenParameters(mIOContext->GetUDPEndPointManager()).SetAddressType(type).SetListenPort(0));
 
-    EXPECT_EQ(err, CHIP_NO_ERROR);
-}
+        EXPECT_EQ(err, CHIP_NO_ERROR);
+    }
 
-void CheckMessageTest(TestContext * ctx, const IPAddress & addr)
-{
-    uint16_t payload_len = sizeof(PAYLOAD);
+    void CheckMessageTest(const IPAddress & addr)
+    {
+        uint16_t payload_len = sizeof(PAYLOAD);
+
+        chip::System::PacketBufferHandle buffer = chip::System::PacketBufferHandle::NewWithData(PAYLOAD, payload_len);
+        EXPECT_FALSE(buffer.IsNull());
 
-    chip::System::PacketBufferHandle buffer = chip::System::PacketBufferHandle::NewWithData(PAYLOAD, payload_len);
-    EXPECT_FALSE(buffer.IsNull());
+        CHIP_ERROR err = CHIP_NO_ERROR;
 
-    CHIP_ERROR err = CHIP_NO_ERROR;
+        Transport::UDP udp;
+
+        err = udp.Init(
+            Transport::UdpListenParameters(mIOContext->GetUDPEndPointManager()).SetAddressType(addr.Type()).SetListenPort(0));
+        EXPECT_EQ(err, CHIP_NO_ERROR);
 
-    Transport::UDP udp;
+        MockTransportMgrDelegate gMockTransportMgrDelegate;
+        TransportMgrBase gTransportMgrBase;
+        gTransportMgrBase.SetSessionManager(&gMockTransportMgrDelegate);
+        gTransportMgrBase.Init(&udp);
 
-    err = udp.Init(Transport::UdpListenParameters(ctx->GetUDPEndPointManager()).SetAddressType(addr.Type()).SetListenPort(0));
-    EXPECT_EQ(err, CHIP_NO_ERROR);
+        ReceiveHandlerCallCount = 0;
 
-    MockTransportMgrDelegate gMockTransportMgrDelegate;
-    TransportMgrBase gTransportMgrBase;
-    gTransportMgrBase.SetSessionManager(&gMockTransportMgrDelegate);
-    gTransportMgrBase.Init(&udp);
+        PacketHeader header;
+        header.SetSourceNodeId(kSourceNodeId).SetDestinationNodeId(kDestinationNodeId).SetMessageCounter(kMessageCounter);
 
-    ReceiveHandlerCallCount = 0;
+        err = header.EncodeBeforeData(buffer);
+        EXPECT_EQ(err, CHIP_NO_ERROR);
 
-    PacketHeader header;
-    header.SetSourceNodeId(kSourceNodeId).SetDestinationNodeId(kDestinationNodeId).SetMessageCounter(kMessageCounter);
+        // Should be able to send a message to itself by just calling send.
+        err = udp.SendMessage(Transport::PeerAddress::UDP(addr, udp.GetBoundPort()), std::move(buffer));
+        EXPECT_EQ(err, CHIP_NO_ERROR);
 
-    err = header.EncodeBeforeData(buffer);
-    EXPECT_EQ(err, CHIP_NO_ERROR);
+        mIOContext->DriveIOUntil(chip::System::Clock::Seconds16(1), []() { return ReceiveHandlerCallCount != 0; });
 
-    // Should be able to send a message to itself by just calling send.
-    err = udp.SendMessage(Transport::PeerAddress::UDP(addr, udp.GetBoundPort()), std::move(buffer));
-    EXPECT_EQ(err, CHIP_NO_ERROR);
+        EXPECT_EQ(ReceiveHandlerCallCount, 1);
+    }
+};
 
-    ctx->DriveIOUntil(chip::System::Clock::Seconds16(1), []() { return ReceiveHandlerCallCount != 0; });
+IOContext * TestUDP::mIOContext = nullptr;
 
-    EXPECT_EQ(ReceiveHandlerCallCount, 1);
-}
+// IPv4 Tests
 
 #if INET_CONFIG_ENABLE_IPV4
 TEST_F(TestUDP, CheckSimpleInitTest4)
 {
-    CheckSimpleInitTest(inContext, IPAddressType::kIPv4);
+    CheckSimpleInitTest(IPAddressType::kIPv4);
 }
 
 TEST_F(TestUDP, CheckMessageTest4)
 {
     IPAddress addr;
     IPAddress::FromString("127.0.0.1", addr);
-    CheckMessageTest(inContext, addr);
+    CheckMessageTest(addr);
 }
 #endif
 
+// IPv6 Tests
+
 TEST_F(TestUDP, CheckSimpleInitTest6)
 {
-    CheckSimpleInitTest(inContext, IPAddressType::kIPv6);
+    CheckSimpleInitTest(IPAddressType::kIPv6);
 }
 
 TEST_F(TestUDP, CheckMessageTest6)
 {
     IPAddress addr;
     IPAddress::FromString("::1", addr);
-    CheckMessageTest(inContext, addr);
+    CheckMessageTest(addr);
 }
+
+} // namespace