From 719000082f74ee340eef22b18ba067d3c9e718ba Mon Sep 17 00:00:00 2001
From: Sean McArthur <sean@seanmonstar.com>
Date: Mon, 29 Jan 2018 15:15:06 -0800
Subject: [PATCH] proxy: fix tcp_with_no_orig_dst test (#229)

Sometimes, the try_read will return a connection error, sometimes it
will just return EOF. Handle both cases.

Closes #226
---
 proxy/tests/transparency.rs | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/proxy/tests/transparency.rs b/proxy/tests/transparency.rs
index 8d0cd503a..2cb4ee19d 100644
--- a/proxy/tests/transparency.rs
+++ b/proxy/tests/transparency.rs
@@ -227,7 +227,9 @@ fn inbound_tcp() {
 fn tcp_with_no_orig_dst() {
     let _ = env_logger::init();
 
-    let srv = server::tcp().run();
+    let srv = server::tcp()
+        .accept(move |_| "don't read me")
+        .run();
     let ctrl = controller::new().run();
     let proxy = proxy::new()
         .controller(ctrl)
@@ -240,7 +242,11 @@ fn tcp_with_no_orig_dst() {
     let tcp_client = client.connect();
     tcp_client.write("custom tcp hello");
 
-    assert!(tcp_client.try_read().is_err());
+    let read = tcp_client
+        .try_read()
+        // This read might be an error, or an empty vec
+        .unwrap_or_else(|_| Vec::new());
+    assert_eq!(read, b"");
 }
 
 #[test]