Add mNumMessagesToAllowBeforeDropping

project-chip · Jul 28, 2022 · fd10374 · fd10374
1 parent 6cd4676
commit fd10374
Showing 1 changed file with 30 additions and 17 deletions.
diff --git a/src/transport/raw/tests/NetworkTestHelpers.h b/src/transport/raw/tests/NetworkTestHelpers.h
@@ -64,7 +64,7 @@ class LoopbackTransportDelegate
 public:
     virtual ~LoopbackTransportDelegate() {}
 
-    // Called by the loopback transport when it drops a message due to a nonzero mNumMessagesToDrop.
+    // Called by the loopback transport when it drops a message due to a nonzero mNumMessagesToDrop/mNumMessagesToAllowBeforeDropping
     virtual void OnMessageDropped() {}
 };
 
@@ -102,19 +102,30 @@ class LoopbackTransport : public Transport::Base
     {
         ReturnErrorOnFailure(mMessageSendError);
         mSentMessageCount++;
-
-        if (mNumMessagesToDrop == 0)
+        bool dropMessage = false;
+        if (mNumMessagesToAllowBeforeDropping > 0)
         {
-            System::PacketBufferHandle receivedMessage = msgBuf.CloneData();
-            mPendingMessageQueue.push(PendingMessageItem(address, std::move(receivedMessage)));
-            mSystemLayer->ScheduleWork(OnMessageReceived, this);
+            --mNumMessagesToAllowBeforeDropping;
         }
-        else
+        else if (mNumMessagesToDrop > 0)
+        {
+            dropMessage = true;
+            --mNumMessagesToDrop;
+        }
+
+        if (dropMessage)
         {
-            mNumMessagesToDrop--;
             mDroppedMessageCount++;
             if (mDelegate != nullptr)
+            {
                 mDelegate->OnMessageDropped();
+            }
+        }
+        else
+        {
+            System::PacketBufferHandle receivedMessage = msgBuf.CloneData();
+            mPendingMessageQueue.push(PendingMessageItem(address, std::move(receivedMessage)));
+            mSystemLayer->ScheduleWork(OnMessageReceived, this);
         }
 
         return CHIP_NO_ERROR;
@@ -124,10 +135,11 @@ class LoopbackTransport : public Transport::Base
 
     void Reset()
     {
-        mNumMessagesToDrop   = 0;
-        mDroppedMessageCount = 0;
-        mSentMessageCount    = 0;
-        mMessageSendError    = CHIP_NO_ERROR;
+        mNumMessagesToDrop                = 0;
+        mDroppedMessageCount              = 0;
+        mSentMessageCount                 = 0;
+        mNumMessagesToAllowBeforeDropping = 0;
+        mMessageSendError                 = CHIP_NO_ERROR;
     }
 
     struct PendingMessageItem
@@ -143,11 +155,12 @@ class LoopbackTransport : public Transport::Base
     System::Layer * mSystemLayer = nullptr;
     std::queue<PendingMessageItem> mPendingMessageQueue;
     Transport::PeerAddress mTxAddress;
-    uint32_t mNumMessagesToDrop           = 0;
-    uint32_t mDroppedMessageCount         = 0;
-    uint32_t mSentMessageCount            = 0;
-    CHIP_ERROR mMessageSendError          = CHIP_NO_ERROR;
-    LoopbackTransportDelegate * mDelegate = nullptr;
+    uint32_t mNumMessagesToDrop                = 0;
+    uint32_t mDroppedMessageCount              = 0;
+    uint32_t mSentMessageCount                 = 0;
+    uint32_t mNumMessagesToAllowBeforeDropping = 0;
+    CHIP_ERROR mMessageSendError               = CHIP_NO_ERROR;
+    LoopbackTransportDelegate * mDelegate      = nullptr;
 };
 
 } // namespace Test