frontend_test.go

Documentation: github.com/jackc/pgconn

     1  package pgconn_test
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"os"
     7  	"testing"
     8  
     9  	"github.com/jackc/pgconn"
    10  	"github.com/jackc/pgproto3/v2"
    11  	"github.com/stretchr/testify/assert"
    12  	"github.com/stretchr/testify/require"
    13  )
    14  
    15  // frontendWrapper allows to hijack a regular frontend, and inject a specific response
    16  type frontendWrapper struct {
    17  	front pgconn.Frontend
    18  
    19  	msg pgproto3.BackendMessage
    20  }
    21  
    22  // frontendWrapper implements the pgconn.Frontend interface
    23  var _ pgconn.Frontend = (*frontendWrapper)(nil)
    24  
    25  func (f *frontendWrapper) Receive() (pgproto3.BackendMessage, error) {
    26  	if f.msg != nil {
    27  		return f.msg, nil
    28  	}
    29  
    30  	return f.front.Receive()
    31  }
    32  
    33  func TestFrontendFatalErrExec(t *testing.T) {
    34  	t.Parallel()
    35  
    36  	config, err := pgconn.ParseConfig(os.Getenv("PGX_TEST_CONN_STRING"))
    37  	require.NoError(t, err)
    38  
    39  	buildFrontend := config.BuildFrontend
    40  	var front *frontendWrapper
    41  
    42  	config.BuildFrontend = func(r io.Reader, w io.Writer) pgconn.Frontend {
    43  		wrapped := buildFrontend(r, w)
    44  		front = &frontendWrapper{wrapped, nil}
    45  
    46  		return front
    47  	}
    48  
    49  	conn, err := pgconn.ConnectConfig(context.Background(), config)
    50  	require.NoError(t, err)
    51  	require.NotNil(t, conn)
    52  	require.NotNil(t, front)
    53  
    54  	// set frontend to return a "FATAL" message on next call
    55  	front.msg = &pgproto3.ErrorResponse{Severity: "FATAL", Message: "unit testing fatal error"}
    56  
    57  	_, err = conn.Exec(context.Background(), "SELECT 1").ReadAll()
    58  	assert.Error(t, err)
    59  
    60  	err = conn.Close(context.Background())
    61  	assert.NoError(t, err)
    62  
    63  	select {
    64  	case <-conn.CleanupDone():
    65  		t.Log("ok, CleanupDone() is not blocking")
    66  
    67  	default:
    68  		assert.Fail(t, "connection closed but CleanupDone() still blocking")
    69  	}
    70  }
    71
View as plain text