Use context timeouts in more tests

Tests should timeout in a reasonable time if something is stuck. In
particular this is important when testing deadlock conditions such as
can occur with the copy protocol if both the client and the server are
blocked writing until the other side does a read.

copy_from_test.go

@@ -19,6 +19,9 @@ func TestConnCopyWithAllQueryExecModes(t *testing.T) {
 		t.Run(mode.String(), func(t *testing.T) {
 			t.Parallel()
 
+			ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+			defer cancel()
+
 			cfg := mustParseConfig(t, os.Getenv("PGX_TEST_DATABASE"))
 			cfg.DefaultQueryExecMode = mode
 			conn := mustConnect(t, cfg)
@@ -39,7 +42,7 @@ func TestConnCopyWithAllQueryExecModes(t *testing.T) {
 				{nil, nil, nil, nil, nil},
 			}
 
-			copyCount, err := conn.CopyFrom(context.Background(), pgx.Identifier{"foo"}, []string{"a", "b", "c", "d", "e"}, pgx.CopyFromRows(inputRows))
+			copyCount, err := conn.CopyFrom(ctx, pgx.Identifier{"foo"}, []string{"a", "b", "c", "d", "e"}, pgx.CopyFromRows(inputRows))
 			if err != nil {
 				t.Errorf("Unexpected error for CopyFrom: %v", err)
 			}
@@ -47,7 +50,7 @@ func TestConnCopyWithAllQueryExecModes(t *testing.T) {
 				t.Errorf("Expected CopyFrom to return %d copied rows, but got %d", len(inputRows), copyCount)
 			}
 
-			rows, err := conn.Query(context.Background(), "select * from foo")
+			rows, err := conn.Query(ctx, "select * from foo")
 			if err != nil {
 				t.Errorf("Unexpected error for Query: %v", err)
 			}
@@ -80,6 +83,9 @@ func TestConnCopyWithKnownOIDQueryExecModes(t *testing.T) {
 		t.Run(mode.String(), func(t *testing.T) {
 			t.Parallel()
 
+			ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+			defer cancel()
+
 			cfg := mustParseConfig(t, os.Getenv("PGX_TEST_DATABASE"))
 			cfg.DefaultQueryExecMode = mode
 			conn := mustConnect(t, cfg)
@@ -102,7 +108,7 @@ func TestConnCopyWithKnownOIDQueryExecModes(t *testing.T) {
 				{nil, nil, nil, nil, nil, nil, nil},
 			}
 
-			copyCount, err := conn.CopyFrom(context.Background(), pgx.Identifier{"foo"}, []string{"a", "b", "c", "d", "e", "f", "g"}, pgx.CopyFromRows(inputRows))
+			copyCount, err := conn.CopyFrom(ctx, pgx.Identifier{"foo"}, []string{"a", "b", "c", "d", "e", "f", "g"}, pgx.CopyFromRows(inputRows))
 			if err != nil {
 				t.Errorf("Unexpected error for CopyFrom: %v", err)
 			}
@@ -110,7 +116,7 @@ func TestConnCopyWithKnownOIDQueryExecModes(t *testing.T) {
 				t.Errorf("Expected CopyFrom to return %d copied rows, but got %d", len(inputRows), copyCount)
 			}
 
-			rows, err := conn.Query(context.Background(), "select * from foo")
+			rows, err := conn.Query(ctx, "select * from foo")
 			if err != nil {
 				t.Errorf("Unexpected error for Query: %v", err)
 			}
@@ -140,6 +146,9 @@ func TestConnCopyWithKnownOIDQueryExecModes(t *testing.T) {
 func TestConnCopyFromSmall(t *testing.T) {
 	t.Parallel()
 
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
 	conn := mustConnectString(t, os.Getenv("PGX_TEST_DATABASE"))
 	defer closeConn(t, conn)
 
@@ -160,7 +169,7 @@ func TestConnCopyFromSmall(t *testing.T) {
 		{nil, nil, nil, nil, nil, nil, nil},
 	}
 
-	copyCount, err := conn.CopyFrom(context.Background(), pgx.Identifier{"foo"}, []string{"a", "b", "c", "d", "e", "f", "g"}, pgx.CopyFromRows(inputRows))
+	copyCount, err := conn.CopyFrom(ctx, pgx.Identifier{"foo"}, []string{"a", "b", "c", "d", "e", "f", "g"}, pgx.CopyFromRows(inputRows))
 	if err != nil {
 		t.Errorf("Unexpected error for CopyFrom: %v", err)
 	}
@@ -168,7 +177,7 @@ func TestConnCopyFromSmall(t *testing.T) {
 		t.Errorf("Expected CopyFrom to return %d copied rows, but got %d", len(inputRows), copyCount)
 	}
 
-	rows, err := conn.Query(context.Background(), "select * from foo")
+	rows, err := conn.Query(ctx, "select * from foo")
 	if err != nil {
 		t.Errorf("Unexpected error for Query: %v", err)
 	}
@@ -196,6 +205,9 @@ func TestConnCopyFromSmall(t *testing.T) {
 func TestConnCopyFromSliceSmall(t *testing.T) {
 	t.Parallel()
 
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
 	conn := mustConnectString(t, os.Getenv("PGX_TEST_DATABASE"))
 	defer closeConn(t, conn)
 
@@ -216,7 +228,7 @@ func TestConnCopyFromSliceSmall(t *testing.T) {
 		{nil, nil, nil, nil, nil, nil, nil},
 	}
 
-	copyCount, err := conn.CopyFrom(context.Background(), pgx.Identifier{"foo"}, []string{"a", "b", "c", "d", "e", "f", "g"},
+	copyCount, err := conn.CopyFrom(ctx, pgx.Identifier{"foo"}, []string{"a", "b", "c", "d", "e", "f", "g"},
 		pgx.CopyFromSlice(len(inputRows), func(i int) ([]any, error) {
 			return inputRows[i], nil
 		}))
@@ -227,7 +239,7 @@ func TestConnCopyFromSliceSmall(t *testing.T) {
 		t.Errorf("Expected CopyFrom to return %d copied rows, but got %d", len(inputRows), copyCount)
 	}
 
-	rows, err := conn.Query(context.Background(), "select * from foo")
+	rows, err := conn.Query(ctx, "select * from foo")
 	if err != nil {
 		t.Errorf("Unexpected error for Query: %v", err)
 	}
@@ -255,6 +267,9 @@ func TestConnCopyFromSliceSmall(t *testing.T) {
 func TestConnCopyFromLarge(t *testing.T) {
 	t.Parallel()
 
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
 	conn := mustConnectString(t, os.Getenv("PGX_TEST_DATABASE"))
 	defer closeConn(t, conn)
 
@@ -279,7 +294,7 @@ func TestConnCopyFromLarge(t *testing.T) {
 		inputRows = append(inputRows, []any{int16(0), int32(1), int64(2), "abc", "efg", time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC), tzedTime, []byte{111, 111, 111, 111}})
 	}
 
-	copyCount, err := conn.CopyFrom(context.Background(), pgx.Identifier{"foo"}, []string{"a", "b", "c", "d", "e", "f", "g", "h"}, pgx.CopyFromRows(inputRows))
+	copyCount, err := conn.CopyFrom(ctx, pgx.Identifier{"foo"}, []string{"a", "b", "c", "d", "e", "f", "g", "h"}, pgx.CopyFromRows(inputRows))
 	if err != nil {
 		t.Errorf("Unexpected error for CopyFrom: %v", err)
 	}
@@ -287,7 +302,7 @@ func TestConnCopyFromLarge(t *testing.T) {
 		t.Errorf("Expected CopyFrom to return %d copied rows, but got %d", len(inputRows), copyCount)
 	}
 
-	rows, err := conn.Query(context.Background(), "select * from foo")
+	rows, err := conn.Query(ctx, "select * from foo")
 	if err != nil {
 		t.Errorf("Unexpected error for Query: %v", err)
 	}
@@ -315,10 +330,12 @@ func TestConnCopyFromLarge(t *testing.T) {
 func TestConnCopyFromEnum(t *testing.T) {
 	t.Parallel()
 
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
 	conn := mustConnectString(t, os.Getenv("PGX_TEST_DATABASE"))
 	defer closeConn(t, conn)
 
-	ctx := context.Background()
 	tx, err := conn.Begin(ctx)
 	require.NoError(t, err)
 	defer tx.Rollback(ctx)
@@ -384,6 +401,9 @@ func TestConnCopyFromEnum(t *testing.T) {
 func TestConnCopyFromJSON(t *testing.T) {
 	t.Parallel()
 
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
 	conn := mustConnectString(t, os.Getenv("PGX_TEST_DATABASE"))
 	defer closeConn(t, conn)
 
@@ -403,7 +423,7 @@ func TestConnCopyFromJSON(t *testing.T) {
 		{nil, nil},
 	}
 
-	copyCount, err := conn.CopyFrom(context.Background(), pgx.Identifier{"foo"}, []string{"a", "b"}, pgx.CopyFromRows(inputRows))
+	copyCount, err := conn.CopyFrom(ctx, pgx.Identifier{"foo"}, []string{"a", "b"}, pgx.CopyFromRows(inputRows))
 	if err != nil {
 		t.Errorf("Unexpected error for CopyFrom: %v", err)
 	}
@@ -411,7 +431,7 @@ func TestConnCopyFromJSON(t *testing.T) {
 		t.Errorf("Expected CopyFrom to return %d copied rows, but got %d", len(inputRows), copyCount)
 	}
 
-	rows, err := conn.Query(context.Background(), "select * from foo")
+	rows, err := conn.Query(ctx, "select * from foo")
 	if err != nil {
 		t.Errorf("Unexpected error for Query: %v", err)
 	}
@@ -461,6 +481,9 @@ func (cfs *clientFailSource) Err() error {
 func TestConnCopyFromFailServerSideMidway(t *testing.T) {
 	t.Parallel()
 
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
 	conn := mustConnectString(t, os.Getenv("PGX_TEST_DATABASE"))
 	defer closeConn(t, conn)
 
@@ -475,7 +498,7 @@ func TestConnCopyFromFailServerSideMidway(t *testing.T) {
 		{int32(3), "def"},
 	}
 
-	copyCount, err := conn.CopyFrom(context.Background(), pgx.Identifier{"foo"}, []string{"a", "b"}, pgx.CopyFromRows(inputRows))
+	copyCount, err := conn.CopyFrom(ctx, pgx.Identifier{"foo"}, []string{"a", "b"}, pgx.CopyFromRows(inputRows))
 	if err == nil {
 		t.Errorf("Expected CopyFrom return error, but it did not")
 	}
@@ -486,7 +509,7 @@ func TestConnCopyFromFailServerSideMidway(t *testing.T) {
 		t.Errorf("Expected CopyFrom to return 0 copied rows, but got %d", copyCount)
 	}
 
-	rows, err := conn.Query(context.Background(), "select * from foo")
+	rows, err := conn.Query(ctx, "select * from foo")
 	if err != nil {
 		t.Errorf("Unexpected error for Query: %v", err)
 	}
@@ -537,6 +560,9 @@ func (fs *failSource) Err() error {
 func TestConnCopyFromFailServerSideMidwayAbortsWithoutWaiting(t *testing.T) {
 	t.Parallel()
 
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
 	conn := mustConnectString(t, os.Getenv("PGX_TEST_DATABASE"))
 	defer closeConn(t, conn)
 
@@ -548,7 +574,7 @@ func TestConnCopyFromFailServerSideMidwayAbortsWithoutWaiting(t *testing.T) {
 
 	startTime := time.Now()
 
-	copyCount, err := conn.CopyFrom(context.Background(), pgx.Identifier{"foo"}, []string{"a"}, &failSource{})
+	copyCount, err := conn.CopyFrom(ctx, pgx.Identifier{"foo"}, []string{"a"}, &failSource{})
 	if err == nil {
 		t.Errorf("Expected CopyFrom return error, but it did not")
 	}
@@ -565,7 +591,7 @@ func TestConnCopyFromFailServerSideMidwayAbortsWithoutWaiting(t *testing.T) {
 		t.Errorf("Failing CopyFrom shouldn't have taken so long: %v", copyTime)
 	}
 
-	rows, err := conn.Query(context.Background(), "select * from foo")
+	rows, err := conn.Query(ctx, "select * from foo")
 	if err != nil {
 		t.Errorf("Unexpected error for Query: %v", err)
 	}
@@ -614,6 +640,9 @@ func (fs *slowFailRaceSource) Err() error {
 func TestConnCopyFromSlowFailRace(t *testing.T) {
 	t.Parallel()
 
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
 	conn := mustConnectString(t, os.Getenv("PGX_TEST_DATABASE"))
 	defer closeConn(t, conn)
 
@@ -622,7 +651,7 @@ func TestConnCopyFromSlowFailRace(t *testing.T) {
 		b bytea not null
 	)`)
 
-	copyCount, err := conn.CopyFrom(context.Background(), pgx.Identifier{"foo"}, []string{"a", "b"}, &slowFailRaceSource{})
+	copyCount, err := conn.CopyFrom(ctx, pgx.Identifier{"foo"}, []string{"a", "b"}, &slowFailRaceSource{})
 	if err == nil {
 		t.Errorf("Expected CopyFrom return error, but it did not")
 	}
@@ -639,6 +668,9 @@ func TestConnCopyFromSlowFailRace(t *testing.T) {
 func TestConnCopyFromCopyFromSourceErrorMidway(t *testing.T) {
 	t.Parallel()
 
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
 	conn := mustConnectString(t, os.Getenv("PGX_TEST_DATABASE"))
 	defer closeConn(t, conn)
 
@@ -646,7 +678,7 @@ func TestConnCopyFromCopyFromSourceErrorMidway(t *testing.T) {
 		a bytea not null
 	)`)
 
-	copyCount, err := conn.CopyFrom(context.Background(), pgx.Identifier{"foo"}, []string{"a"}, &clientFailSource{})
+	copyCount, err := conn.CopyFrom(ctx, pgx.Identifier{"foo"}, []string{"a"}, &clientFailSource{})
 	if err == nil {
 		t.Errorf("Expected CopyFrom return error, but it did not")
 	}
@@ -654,7 +686,7 @@ func TestConnCopyFromCopyFromSourceErrorMidway(t *testing.T) {
 		t.Errorf("Expected CopyFrom to return 0 copied rows, but got %d", copyCount)
 	}
 
-	rows, err := conn.Query(context.Background(), "select * from foo")
+	rows, err := conn.Query(ctx, "select * from foo")
 	if err != nil {
 		t.Errorf("Unexpected error for Query: %v", err)
 	}
@@ -699,6 +731,9 @@ func (cfs *clientFinalErrSource) Err() error {
 func TestConnCopyFromCopyFromSourceErrorEnd(t *testing.T) {
 	t.Parallel()
 
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
 	conn := mustConnectString(t, os.Getenv("PGX_TEST_DATABASE"))
 	defer closeConn(t, conn)
 
@@ -706,7 +741,7 @@ func TestConnCopyFromCopyFromSourceErrorEnd(t *testing.T) {
 		a bytea not null
 	)`)
 
-	copyCount, err := conn.CopyFrom(context.Background(), pgx.Identifier{"foo"}, []string{"a"}, &clientFinalErrSource{})
+	copyCount, err := conn.CopyFrom(ctx, pgx.Identifier{"foo"}, []string{"a"}, &clientFinalErrSource{})
 	if err == nil {
 		t.Errorf("Expected CopyFrom return error, but it did not")
 	}
@@ -714,7 +749,7 @@ func TestConnCopyFromCopyFromSourceErrorEnd(t *testing.T) {
 		t.Errorf("Expected CopyFrom to return 0 copied rows, but got %d", copyCount)
 	}
 
-	rows, err := conn.Query(context.Background(), "select * from foo")
+	rows, err := conn.Query(ctx, "select * from foo")
 	if err != nil {
 		t.Errorf("Unexpected error for Query: %v", err)
 	}
@@ -742,6 +777,9 @@ func TestConnCopyFromCopyFromSourceErrorEnd(t *testing.T) {
 func TestConnCopyFromAutomaticStringConversion(t *testing.T) {
 	t.Parallel()
 
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
 	conn := mustConnectString(t, os.Getenv("PGX_TEST_DATABASE"))
 	defer closeConn(t, conn)
 
@@ -755,11 +793,11 @@ func TestConnCopyFromAutomaticStringConversion(t *testing.T) {
 		{8},
 	}
 
-	copyCount, err := conn.CopyFrom(context.Background(), pgx.Identifier{"foo"}, []string{"a"}, pgx.CopyFromRows(inputRows))
+	copyCount, err := conn.CopyFrom(ctx, pgx.Identifier{"foo"}, []string{"a"}, pgx.CopyFromRows(inputRows))
 	require.NoError(t, err)
 	require.EqualValues(t, len(inputRows), copyCount)
 
-	rows, _ := conn.Query(context.Background(), "select * from foo")
+	rows, _ := conn.Query(ctx, "select * from foo")
 	nums, err := pgx.CollectRows(rows, pgx.RowTo[int64])
 	require.NoError(t, err)