From a9e7e3acbc04145211116a11959c4db176a5df9a Mon Sep 17 00:00:00 2001
From: Jack Christensen <jack@jackchristensen.com>
Date: Sat, 4 Feb 2017 16:03:20 -0600
Subject: [PATCH] Extract connection dead on server test

---
 conn_test.go   | 11 +----------
 helper_test.go | 17 ++++++++++++++++-
 query_test.go  | 18 ++----------------
 3 files changed, 19 insertions(+), 27 deletions(-)

diff --git a/conn_test.go b/conn_test.go
index a9cf02c9..e92c7ca3 100644
--- a/conn_test.go
+++ b/conn_test.go
@@ -872,16 +872,7 @@ func TestExecContextCancelationCancelsQuery(t *testing.T) {
 		t.Fatal("Expected context.Canceled err, got %v", err)
 	}
 
-	time.Sleep(500 * time.Millisecond)
-
-	checkConn := mustConnect(t, *defaultConnConfig)
-	defer closeConn(t, checkConn)
-
-	var found bool
-	err = checkConn.QueryRow("select true from pg_stat_activity where pid=$1", conn.Pid).Scan(&found)
-	if err != pgx.ErrNoRows {
-		t.Fatal("Expected context canceled connection to be disconnected from server, but it wasn't")
-	}
+	ensureConnDeadOnServer(t, conn, *defaultConnConfig)
 }
 
 func TestPrepare(t *testing.T) {
diff --git a/helper_test.go b/helper_test.go
index eff731e8..997ae26f 100644
--- a/helper_test.go
+++ b/helper_test.go
@@ -21,7 +21,6 @@ func mustReplicationConnect(t testing.TB, config pgx.ConnConfig) *pgx.Replicatio
 	return conn
 }
 
-
 func closeConn(t testing.TB, conn *pgx.Conn) {
 	err := conn.Close()
 	if err != nil {
@@ -72,3 +71,19 @@ func ensureConnValid(t *testing.T, conn *pgx.Conn) {
 		t.Error("Wrong values returned")
 	}
 }
+
+func ensureConnDeadOnServer(t *testing.T, conn *pgx.Conn, config pgx.ConnConfig) {
+	checkConn := mustConnect(t, config)
+	defer closeConn(t, checkConn)
+
+	for i := 0; i < 10; i++ {
+		var found bool
+		err := checkConn.QueryRow("select true from pg_stat_activity where pid=$1", conn.Pid).Scan(&found)
+		if err == pgx.ErrNoRows {
+			return
+		} else if err != nil {
+			t.Fatalf("Unable to check if conn is dead on server: %v", err)
+		}
+	}
+	t.Fatal("Expected conn to be disconnected from server, but it wasn't")
+}
diff --git a/query_test.go b/query_test.go
index 6909ba1e..40886f2e 100644
--- a/query_test.go
+++ b/query_test.go
@@ -1513,14 +1513,7 @@ func TestQueryContextCancelationCancelsQuery(t *testing.T) {
 		t.Fatal("Expected context.Canceled error, got %v", rows.Err())
 	}
 
-	checkConn := mustConnect(t, *defaultConnConfig)
-	defer closeConn(t, checkConn)
-
-	var found bool
-	err = checkConn.QueryRow("select true from pg_stat_activity where pid=$1", conn.Pid).Scan(&found)
-	if err != pgx.ErrNoRows {
-		t.Fatal("Expected context canceled connection to be disconnected from server, but it wasn't")
-	}
+	ensureConnDeadOnServer(t, conn, *defaultConnConfig)
 }
 
 func TestQueryRowContextSuccess(t *testing.T) {
@@ -1580,12 +1573,5 @@ func TestQueryRowContextCancelationCancelsQuery(t *testing.T) {
 		t.Fatal("Expected context.Canceled error, got %v", err)
 	}
 
-	checkConn := mustConnect(t, *defaultConnConfig)
-	defer closeConn(t, checkConn)
-
-	var found bool
-	err = checkConn.QueryRow("select true from pg_stat_activity where pid=$1", conn.Pid).Scan(&found)
-	if err != pgx.ErrNoRows {
-		t.Fatal("Expected context canceled connection to be disconnected from server, but it wasn't")
-	}
+	ensureConnDeadOnServer(t, conn, *defaultConnConfig)
 }