github.com/jackc/pgx/v5@v5.5.5/internal/pgmock/pgmock_test.go (about)

     1  package pgmock_test
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"net"
     7  	"strings"
     8  	"testing"
     9  	"time"
    10  
    11  	"github.com/jackc/pgx/v5/internal/pgmock"
    12  	"github.com/jackc/pgx/v5/pgconn"
    13  	"github.com/jackc/pgx/v5/pgproto3"
    14  
    15  	"github.com/stretchr/testify/assert"
    16  	"github.com/stretchr/testify/require"
    17  )
    18  
    19  func TestScript(t *testing.T) {
    20  	script := &pgmock.Script{
    21  		Steps: pgmock.AcceptUnauthenticatedConnRequestSteps(),
    22  	}
    23  	script.Steps = append(script.Steps, pgmock.ExpectMessage(&pgproto3.Query{String: "select 42"}))
    24  	script.Steps = append(script.Steps, pgmock.SendMessage(&pgproto3.RowDescription{
    25  		Fields: []pgproto3.FieldDescription{
    26  			{
    27  				Name:                 []byte("?column?"),
    28  				TableOID:             0,
    29  				TableAttributeNumber: 0,
    30  				DataTypeOID:          23,
    31  				DataTypeSize:         4,
    32  				TypeModifier:         -1,
    33  				Format:               0,
    34  			},
    35  		},
    36  	}))
    37  	script.Steps = append(script.Steps, pgmock.SendMessage(&pgproto3.DataRow{
    38  		Values: [][]byte{[]byte("42")},
    39  	}))
    40  	script.Steps = append(script.Steps, pgmock.SendMessage(&pgproto3.CommandComplete{CommandTag: []byte("SELECT 1")}))
    41  	script.Steps = append(script.Steps, pgmock.SendMessage(&pgproto3.ReadyForQuery{TxStatus: 'I'}))
    42  	script.Steps = append(script.Steps, pgmock.ExpectMessage(&pgproto3.Terminate{}))
    43  
    44  	ln, err := net.Listen("tcp", "127.0.0.1:")
    45  	require.NoError(t, err)
    46  	defer ln.Close()
    47  
    48  	serverErrChan := make(chan error, 1)
    49  	go func() {
    50  		defer close(serverErrChan)
    51  
    52  		conn, err := ln.Accept()
    53  		if err != nil {
    54  			serverErrChan <- err
    55  			return
    56  		}
    57  		defer conn.Close()
    58  
    59  		err = conn.SetDeadline(time.Now().Add(time.Second))
    60  		if err != nil {
    61  			serverErrChan <- err
    62  			return
    63  		}
    64  
    65  		err = script.Run(pgproto3.NewBackend(conn, conn))
    66  		if err != nil {
    67  			serverErrChan <- err
    68  			return
    69  		}
    70  	}()
    71  
    72  	host, port, _ := strings.Cut(ln.Addr().String(), ":")
    73  	connStr := fmt.Sprintf("sslmode=disable host=%s port=%s", host, port)
    74  
    75  	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
    76  	defer cancel()
    77  	pgConn, err := pgconn.Connect(ctx, connStr)
    78  	require.NoError(t, err)
    79  	results, err := pgConn.Exec(ctx, "select 42").ReadAll()
    80  	assert.NoError(t, err)
    81  
    82  	assert.Len(t, results, 1)
    83  	assert.Nil(t, results[0].Err)
    84  	assert.Equal(t, "SELECT 1", results[0].CommandTag.String())
    85  	assert.Len(t, results[0].Rows, 1)
    86  	assert.Equal(t, "42", string(results[0].Rows[0][0]))
    87  
    88  	pgConn.Close(ctx)
    89  
    90  	assert.NoError(t, <-serverErrChan)
    91  }