Make test harness connection shutdown a little cleaner

vectordotdev · Dec 11, 2018 · 39bd742 · 39bd742
1 parent 282f4c9
commit 39bd742
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 8 deletions.
diff --git a/benches/bench.rs b/benches/bench.rs
@@ -548,11 +548,9 @@ fn send_lines(
                 .forward(out)
                 .map(|(_source, sink)| sink)
                 .and_then(|sink| {
-                    // This waits for FIN from the server so we don't start shutting it down before it's fully received the test data
                     let socket = sink.into_inner().into_inner();
-                    socket.shutdown(std::net::Shutdown::Write).unwrap();
-                    FramedRead::new(socket, BytesCodec::new())
-                        .for_each(|_| Ok(()))
+                    tokio::io::shutdown(socket)
+                        .map(|_| ())
                         .map_err(|e| panic!("{:}", e))
                 })
         })

diff --git a/tests/basic_tcp.rs b/tests/basic_tcp.rs
@@ -451,11 +451,9 @@ fn send_lines(
                 .forward(out)
                 .map(|(_source, sink)| sink)
                 .and_then(|sink| {
-                    // This waits for FIN from the server so we don't start shutting it down before it's fully received the test data
                     let socket = sink.into_inner().into_inner();
-                    socket.shutdown(std::net::Shutdown::Write).unwrap();
-                    FramedRead::new(socket, BytesCodec::new())
-                        .for_each(|_| Ok(()))
+                    tokio::io::shutdown(socket)
+                        .map(|_| ())
                         .map_err(|e| panic!("{:}", e))
                 })
         })