1 files changed, 18 insertions, 8 deletions
diff --git a/tests/unit_tests/epee_levin_protocol_handler_async.cpp b/tests/unit_tests/epee_levin_protocol_handler_async.cpp
index a499fa608..d9de99b8b 100644
--- a/tests/unit_tests/epee_levin_protocol_handler_async.cpp
+++ b/tests/unit_tests/epee_levin_protocol_handler_async.cpp
@@ -59,13 +59,13 @@ namespace
     {
     }
 
-    virtual int invoke(int command, const epee::span<const uint8_t> in_buff, epee::byte_slice& buff_out, test_levin_connection_context& context)
+    virtual int invoke(int command, const epee::span<const uint8_t> in_buff, epee::byte_stream& buff_out, test_levin_connection_context& context)
     {
       m_invoke_counter.inc();
       boost::unique_lock<boost::mutex> lock(m_mutex);
       m_last_command = command;
       m_last_in_buf = std::string((const char*)in_buff.data(), in_buff.size());
-      buff_out = m_invoke_out_buf.clone();
+      buff_out.write(epee::to_span(m_invoke_out_buf));
       return m_return_code;
     }
 
@@ -434,8 +434,11 @@ TEST_F(positive_test_connection_to_levin_protocol_handler_calls, handler_process
   const int expected_command = 4673261;
   const std::string in_data(256, 'e');
 
+  epee::levin::message_writer message{};
+  message.buffer.write(epee::to_span(in_data));
+
   const epee::byte_slice noise = epee::levin::make_noise_notify(1024);
-  const epee::byte_slice notify = epee::levin::make_notify(expected_command, epee::strspan<std::uint8_t>(in_data));
+  const epee::byte_slice notify = message.finalize_notify(expected_command);
 
   test_connection_ptr conn = create_connection();
 
@@ -468,11 +471,16 @@ TEST_F(positive_test_connection_to_levin_protocol_handler_calls, handler_process
   const int expected_command = 4673261;
   const int expected_fragmented_command = 46732;
   const std::string in_data(256, 'e');
-  std::string in_fragmented_data(1024 * 4, 'c');
+
+  epee::levin::message_writer message{};
+  message.buffer.write(epee::to_span(in_data));
+
+  epee::levin::message_writer in_fragmented_data;
+  in_fragmented_data.buffer.put_n('c', 1024 * 4);
 
   const epee::byte_slice noise = epee::levin::make_noise_notify(1024);
-  const epee::byte_slice notify = epee::levin::make_notify(expected_command, epee::strspan<std::uint8_t>(in_data));
-  epee::byte_slice fragmented = epee::levin::make_fragmented_notify(noise, expected_fragmented_command, epee::strspan<std::uint8_t>(in_fragmented_data));
+  const epee::byte_slice notify = message.finalize_notify(expected_command);
+  epee::byte_slice fragmented = epee::levin::make_fragmented_notify(noise.size(), expected_fragmented_command, std::move(in_fragmented_data));
 
   EXPECT_EQ(5u, fragmented.size() / 1024);
   EXPECT_EQ(0u, fragmented.size() % 1024);
@@ -497,11 +505,13 @@ TEST_F(positive_test_connection_to_levin_protocol_handler_calls, handler_process
     ASSERT_TRUE(conn->m_protocol_handler.handle_recv(next.data(), next.size()));
   }
 
-  in_fragmented_data.resize(((1024 - sizeof(epee::levin::bucket_head2)) * 5) - sizeof(epee::levin::bucket_head2)); // add padding zeroes
+  std::string compare_buffer(1024 * 4, 'c');
+  compare_buffer.resize(((1024 - sizeof(epee::levin::bucket_head2)) * 5) - sizeof(epee::levin::bucket_head2)); // add padding zeroes
+
   ASSERT_EQ(4u, m_commands_handler.notify_counter());
   ASSERT_EQ(0u, m_commands_handler.invoke_counter());
   ASSERT_EQ(expected_fragmented_command, m_commands_handler.last_command());
-  ASSERT_EQ(in_fragmented_data, m_commands_handler.last_in_buf());
+  ASSERT_EQ(compare_buffer, m_commands_handler.last_in_buf());
   ASSERT_EQ(0u, conn->send_counter());
   ASSERT_TRUE(conn->last_send_data().empty());