From 2b9033cf1697adb7bf067e7abb3bb1ef546c64d9 Mon Sep 17 00:00:00 2001 From: Sean McArthur Date: Tue, 10 Apr 2018 19:47:00 -0700 Subject: [PATCH] proxy: fix flaky tcp graceful shutdown test (#735) --- proxy/tests/shutdown.rs | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/proxy/tests/shutdown.rs b/proxy/tests/shutdown.rs index 2582ee4d7..7fc726600 100644 --- a/proxy/tests/shutdown.rs +++ b/proxy/tests/shutdown.rs @@ -131,9 +131,17 @@ fn tcp_waits_for_proxies_to_close() { let msg2 = "custom tcp bye"; let srv = server::tcp() - .accept(move |read| { - assert_eq!(read, msg1.as_bytes()); - msg2 + // Trigger a shutdown while TCP stream is busy + .accept_fut(move |sock| { + shdn.signal(); + tokio_io::io::read(sock, vec![0; 256]) + .and_then(move |(sock, vec, n)| { + assert_eq!(&vec[..n], msg1.as_bytes()); + + tokio_io::io::write_all(sock, msg2.as_bytes()) + }) + .map(|_| ()) + .map_err(|e| panic!("tcp server error: {}", e)) }) .run(); let ctrl = controller::new().run(); @@ -147,8 +155,6 @@ fn tcp_waits_for_proxies_to_close() { let tcp_client = client.connect(); - shdn.signal(); - tcp_client.write(msg1); assert_eq!(tcp_client.read(), msg2.as_bytes()); }