diff --git a/src/io/rsm/raft.hpp b/src/io/rsm/raft.hpp
index dc1a1e9b3..07b6b8a4b 100644
--- a/src/io/rsm/raft.hpp
+++ b/src/io/rsm/raft.hpp
@@ -156,14 +156,17 @@ struct Follower {
 using Role = std::variant<Candidate, Leader, Follower>;
 
 /*
+
 all ReplicatedState classes should have an Apply method
-that returns our WriteResponseValue:
+that returns our WriteResponseValue after consensus, and
+a Read method that returns our ReadResponseValue without
+requiring consensus.
 
 ReadResponse Read(ReadOperation);
 WriteResponseValue ReplicatedState::Apply(WriteRequest);
 
-for examples:
-if the state is uint64_t, and WriteRequest is `struct PlusOne {};`,
+For example:
+If the state is uint64_t, and WriteRequest is `struct PlusOne {};`,
 and WriteResponseValue is also uint64_t (the new value), then
 each call to state.Apply(PlusOne{}) will return the new value
 after incrementing it. 0, 1, 2, 3... and this will be sent back
diff --git a/tests/simulation/sharded_map.cpp b/tests/simulation/sharded_map.cpp
index 8b0928ecc..e990e2045 100644
--- a/tests/simulation/sharded_map.cpp
+++ b/tests/simulation/sharded_map.cpp
@@ -58,6 +58,7 @@ using memgraph::io::simulator::SimulatorTransport;
 using StorageClient =
     RsmClient<Io<SimulatorTransport>, StorageWriteRequest, StorageWriteResponse, StorageGetRequest, StorageGetResponse>;
 namespace {
+
 ShardMap CreateDummyShardmap(memgraph::coordinator::Address a_io_1, memgraph::coordinator::Address a_io_2,
                              memgraph::coordinator::Address a_io_3, memgraph::coordinator::Address b_io_1,
                              memgraph::coordinator::Address b_io_2, memgraph::coordinator::Address b_io_3) {
@@ -193,8 +194,6 @@ int main() {
   auto b_thread_3 = std::jthread(RunStorageRaft<SimulatorTransport>, std::move(b_3));
   simulator.IncrementServerCountAndWaitForQuiescentState(b_addrs[2]);
 
-  std::cout << "beginning test after servers have become quiescent" << std::endl;
-
   // Spin up coordinators
 
   Io<SimulatorTransport> c_io_1 = simulator.RegisterNew();
@@ -220,6 +219,8 @@ int main() {
   auto c_thread_3 = std::jthread([c_3]() mutable { c_3.Run(); });
   simulator.IncrementServerCountAndWaitForQuiescentState(c_addrs[2]);
 
+  std::cout << "beginning test after servers have become quiescent" << std::endl;
+
   // Have client contact coordinator RSM for a new transaction ID and
   // also get the current shard map
   using CoordinatorClient =
@@ -259,7 +260,9 @@ int main() {
     // Transaction ID to be used later...
     auto transaction_id = res.new_hlc;
 
-    client_shard_map = res.fresher_shard_map.value();
+    if (res.fresher_shard_map) {
+      client_shard_map = res.fresher_shard_map.value();
+    }
 
     // TODO(gabor) check somewhere in the call chain if the entries are actually valid
     // for (auto &[key, val] : client_shard_map.GetShards()) {
diff --git a/tests/unit/CMakeLists.txt b/tests/unit/CMakeLists.txt
index 9de4860ef..3050ed24b 100644
--- a/tests/unit/CMakeLists.txt
+++ b/tests/unit/CMakeLists.txt
@@ -391,12 +391,6 @@ add_custom_target(test_lcp ALL DEPENDS ${CMAKE_CURRENT_BINARY_DIR}/test_lcp)
 add_test(test_lcp ${CMAKE_CURRENT_BINARY_DIR}/test_lcp)
 add_dependencies(memgraph__unit test_lcp)
 
-# Test websocket
-find_package(Boost REQUIRED)
-
-add_unit_test(websocket.cpp)
-target_link_libraries(${test_prefix}websocket mg-communication Boost::headers)
-
 # Test future
 add_unit_test(future.cpp)
-target_link_libraries(${test_prefix}future mg-io)
\ No newline at end of file
+target_link_libraries(${test_prefix}future mg-io)