1 files changed, 29 insertions, 11 deletions
diff --git a/host/tests/common/mock_zero_copy.cpp b/host/tests/common/mock_zero_copy.cpp
index eee96234d..bc49c3f10 100644
--- a/host/tests/common/mock_zero_copy.cpp
+++ b/host/tests/common/mock_zero_copy.cpp
@@ -11,8 +11,12 @@
 using namespace uhd::transport;
 
 mock_zero_copy::mock_zero_copy(
-    vrt::if_packet_info_t::link_type_t link_type
-) : _link_type(link_type) {
+    vrt::if_packet_info_t::link_type_t link_type,
+    size_t recv_frame_size,
+    size_t send_frame_size
+) : _link_type(link_type)
+  , _recv_frame_size(recv_frame_size)
+  , _send_frame_size(send_frame_size) {
 }
 
 uhd::transport::managed_recv_buffer::sptr mock_zero_copy::get_recv_buff(double) {
@@ -22,18 +26,32 @@ uhd::transport::managed_recv_buffer::sptr mock_zero_copy::get_recv_buff(double)
     if (_rx_mems.empty()) {
         return uhd::transport::managed_recv_buffer::sptr(); // timeout
     }
-    _mrbs.push_back(boost::make_shared<mock_mrb>());
+
     uhd::transport::managed_recv_buffer::sptr mrb =
-        _mrbs.back()->get_new(_rx_mems.front(), _rx_lens.front());
-    _rx_mems.pop_front();
-    _rx_lens.pop_front();
+        _mrb.get_new(_rx_mems.front(), _rx_lens.front());
+
+    if (not _reuse_recv_memory) {
+        _rx_mems.pop_front();
+        _rx_lens.pop_front();
+    }
+
     return mrb;
 }
 
 uhd::transport::managed_send_buffer::sptr mock_zero_copy::get_send_buff(double) {
-    _msbs.push_back(boost::make_shared<mock_msb>());
-    _tx_mems.push_back(
-        boost::shared_array<uint8_t>(new uint8_t[SEND_BUFF_SIZE]));
-    _tx_lens.push_back(SEND_BUFF_SIZE);
-    return _msbs.back()->get_new(_tx_mems.back(), &_tx_lens.back());
+    if (not _reuse_send_memory or _tx_mems.size() == 0) {
+        _tx_mems.push_back(
+            boost::shared_array<uint8_t>(new uint8_t[_send_frame_size]));
+        _tx_lens.push_back(_send_frame_size);
+    }
+
+    return _msb.get_new(_tx_mems.back(), &_tx_lens.back());
+}
+
+void mock_zero_copy::set_reuse_recv_memory(bool reuse_recv) {
+    _reuse_recv_memory = reuse_recv;
+}
+
+void mock_zero_copy::set_reuse_send_memory(bool reuse_send) {
+    _reuse_send_memory = reuse_send;
 }