From b61ca0cae086ad91db76754e2ae7254230017068 Mon Sep 17 00:00:00 2001
From: Eno Compton <enocom@google.com>
Date: Wed, 22 Nov 2023 11:16:49 -0700
Subject: [PATCH] Improve test reliability

---
 cmd/wait_test.go | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/cmd/wait_test.go b/cmd/wait_test.go
index 94af153fb..fcc332b67 100644
--- a/cmd/wait_test.go
+++ b/cmd/wait_test.go
@@ -15,8 +15,10 @@
 package cmd
 
 import (
+	"io"
 	"net"
 	"testing"
+	"time"
 )
 
 func TestWaitCommandFlags(t *testing.T) {
@@ -35,6 +37,11 @@ func TestWaitCommandFlags(t *testing.T) {
 			return
 		}
 		defer conn.Close()
+		// Use a read deadline to produce read error
+		conn.SetReadDeadline(time.Now().Add(100 * time.Millisecond))
+		// Read client request first.
+		io.ReadAll(conn)
+		// Write a generic 200 response back.
 		conn.Write([]byte("HTTP/1.1 200 OK\r\n\r\n"))
 	}()
 
@@ -53,7 +60,7 @@ func TestWaitCommandFails(t *testing.T) {
 	_, err := invokeProxyCommand([]string{
 		"wait",
 		// assuming default host and port
-		"--max=500ms",
+		"--max=100ms",
 	})
 	if err == nil {
 		t.Fatal("wait should fail when endpoint does not respond")