github.com/parquet-go/parquet-go@v0.21.1-0.20240501160520-b3c3a0c3ed6f/buffer_pool_test.go (about)

     1  package parquet_test
     2  
     3  import (
     4  	"bytes"
     5  	"io"
     6  	"strings"
     7  	"testing"
     8  	"testing/iotest"
     9  
    10  	"github.com/parquet-go/parquet-go"
    11  )
    12  
    13  func TestBufferPool(t *testing.T) {
    14  	testBufferPool(t, parquet.NewBufferPool())
    15  }
    16  
    17  func TestFileBufferPool(t *testing.T) {
    18  	testBufferPool(t, parquet.NewFileBufferPool("/tmp", "buffers.*"))
    19  }
    20  
    21  func testBufferPool(t *testing.T, pool parquet.BufferPool) {
    22  	tests := []struct {
    23  		scenario string
    24  		function func(*testing.T, parquet.BufferPool)
    25  	}{
    26  		{
    27  			scenario: "write bytes",
    28  			function: testBufferPoolWriteBytes,
    29  		},
    30  
    31  		{
    32  			scenario: "write string",
    33  			function: testBufferPoolWriteString,
    34  		},
    35  
    36  		{
    37  			scenario: "copy to buffer",
    38  			function: testBufferPoolCopyToBuffer,
    39  		},
    40  
    41  		{
    42  			scenario: "copy from buffer",
    43  			function: testBufferPoolCopyFromBuffer,
    44  		},
    45  	}
    46  
    47  	for _, test := range tests {
    48  		t.Run(test.scenario, func(t *testing.T) { test.function(t, pool) })
    49  	}
    50  }
    51  
    52  func testBufferPoolWriteBytes(t *testing.T, pool parquet.BufferPool) {
    53  	const content = "Hello World!"
    54  
    55  	buffer := pool.GetBuffer()
    56  	defer pool.PutBuffer(buffer)
    57  
    58  	_, err := buffer.Write([]byte(content))
    59  	if err != nil {
    60  		t.Fatal(err)
    61  	}
    62  	assertBufferContent(t, buffer, content)
    63  }
    64  
    65  func testBufferPoolWriteString(t *testing.T, pool parquet.BufferPool) {
    66  	const content = "Hello World!"
    67  
    68  	buffer := pool.GetBuffer()
    69  	defer pool.PutBuffer(buffer)
    70  
    71  	_, err := io.WriteString(buffer, content)
    72  	if err != nil {
    73  		t.Fatal(err)
    74  	}
    75  
    76  	assertBufferContent(t, buffer, content)
    77  }
    78  
    79  func testBufferPoolCopyToBuffer(t *testing.T, pool parquet.BufferPool) {
    80  	const content = "ABC"
    81  
    82  	buffer := pool.GetBuffer()
    83  	defer pool.PutBuffer(buffer)
    84  
    85  	reader := strings.NewReader(content)
    86  	_, err := io.Copy(buffer, struct{ io.Reader }{reader})
    87  	if err != nil {
    88  		t.Fatal(err)
    89  	}
    90  
    91  	assertBufferContent(t, buffer, content)
    92  }
    93  
    94  func testBufferPoolCopyFromBuffer(t *testing.T, pool parquet.BufferPool) {
    95  	const content = "0123456789"
    96  
    97  	buffer := pool.GetBuffer()
    98  	defer pool.PutBuffer(buffer)
    99  
   100  	if _, err := io.WriteString(buffer, content); err != nil {
   101  		t.Fatal(err)
   102  	}
   103  	if _, err := buffer.Seek(0, io.SeekStart); err != nil {
   104  		t.Fatal(err)
   105  	}
   106  
   107  	writer := new(bytes.Buffer)
   108  	_, err := io.Copy(struct{ io.Writer }{writer}, buffer)
   109  	if err != nil {
   110  		t.Fatal(err)
   111  	}
   112  
   113  	assertBufferContent(t, bytes.NewReader(writer.Bytes()), content)
   114  }
   115  
   116  func assertBufferContent(t *testing.T, b io.ReadSeeker, s string) {
   117  	t.Helper()
   118  
   119  	offset, err := b.Seek(0, io.SeekStart)
   120  	if err != nil {
   121  		t.Error("seek:", err)
   122  	}
   123  	if offset != 0 {
   124  		t.Errorf("seek: invalid offset returned: want=0 got=%d", offset)
   125  	}
   126  	if err := iotest.TestReader(b, []byte(s)); err != nil {
   127  		t.Error("iotest:", err)
   128  	}
   129  }