github.com/jackc/pgx/v5@v5.5.5/pgproto3/chunkreader_test.go (about)

     1  package pgproto3
     2  
     3  import (
     4  	"bytes"
     5  	"math/rand"
     6  	"testing"
     7  )
     8  
     9  func TestChunkReaderNextDoesNotReadIfAlreadyBuffered(t *testing.T) {
    10  	server := &bytes.Buffer{}
    11  	r := newChunkReader(server, 4)
    12  
    13  	src := []byte{1, 2, 3, 4}
    14  	server.Write(src)
    15  
    16  	n1, err := r.Next(2)
    17  	if err != nil {
    18  		t.Fatal(err)
    19  	}
    20  	if !bytes.Equal(n1, src[0:2]) {
    21  		t.Fatalf("Expected read bytes to be %v, but they were %v", src[0:2], n1)
    22  	}
    23  
    24  	n2, err := r.Next(2)
    25  	if err != nil {
    26  		t.Fatal(err)
    27  	}
    28  	if !bytes.Equal(n2, src[2:4]) {
    29  		t.Fatalf("Expected read bytes to be %v, but they were %v", src[2:4], n2)
    30  	}
    31  
    32  	if !bytes.Equal((*r.buf)[:len(src)], src) {
    33  		t.Fatalf("Expected r.buf to be %v, but it was %v", src, r.buf)
    34  	}
    35  
    36  	_, err = r.Next(0) // Trigger the buffer reset.
    37  	if err != nil {
    38  		t.Fatal(err)
    39  	}
    40  
    41  	if r.rp != 0 {
    42  		t.Fatalf("Expected r.rp to be %v, but it was %v", 0, r.rp)
    43  	}
    44  	if r.wp != 0 {
    45  		t.Fatalf("Expected r.wp to be %v, but it was %v", 0, r.wp)
    46  	}
    47  }
    48  
    49  type randomReader struct {
    50  	rnd *rand.Rand
    51  }
    52  
    53  // Read reads a random number of random bytes.
    54  func (r *randomReader) Read(p []byte) (n int, err error) {
    55  	n = r.rnd.Intn(len(p) + 1)
    56  	return r.rnd.Read(p[:n])
    57  }
    58  
    59  func TestChunkReaderNextFuzz(t *testing.T) {
    60  	rr := &randomReader{rnd: rand.New(rand.NewSource(1))}
    61  	r := newChunkReader(rr, 8192)
    62  
    63  	randomSizes := rand.New(rand.NewSource(0))
    64  
    65  	for i := 0; i < 100000; i++ {
    66  		size := randomSizes.Intn(16384) + 1
    67  		buf, err := r.Next(size)
    68  		if err != nil {
    69  			t.Fatal(err)
    70  		}
    71  		if len(buf) != size {
    72  			t.Fatalf("Expected to get %v bytes but got %v bytes", size, len(buf))
    73  		}
    74  	}
    75  }