github.com/vc42/parquet-go@v0.0.0-20240320194221-1a9adb5f23f5/page_buffer_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/vc42/parquet-go"
    11  )
    12  
    13  func TestPageBufferPool(t *testing.T) {
    14  	testPageBufferPool(t, parquet.NewPageBufferPool())
    15  }
    16  
    17  func TestFileBufferPool(t *testing.T) {
    18  	testPageBufferPool(t, parquet.NewFileBufferPool("/tmp", "buffers.*"))
    19  }
    20  
    21  func testPageBufferPool(t *testing.T, pool parquet.PageBufferPool) {
    22  	tests := []struct {
    23  		scenario string
    24  		function func(*testing.T, parquet.PageBufferPool)
    25  	}{
    26  		{
    27  			scenario: "write bytes",
    28  			function: testPageBufferPoolWriteBytes,
    29  		},
    30  
    31  		{
    32  			scenario: "write string",
    33  			function: testPageBufferPoolWriteString,
    34  		},
    35  
    36  		{
    37  			scenario: "copy to buffer",
    38  			function: testPageBufferPoolCopyToBuffer,
    39  		},
    40  
    41  		{
    42  			scenario: "copy from buffer",
    43  			function: testPageBufferPoolCopyFromBuffer,
    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 testPageBufferPoolWriteBytes(t *testing.T, pool parquet.PageBufferPool) {
    53  	const content = "Hello World!"
    54  
    55  	buffer := pool.GetPageBuffer()
    56  	_, err := buffer.Write([]byte(content))
    57  	if err != nil {
    58  		t.Fatal(err)
    59  	}
    60  
    61  	assertBufferContent(t, buffer, content)
    62  }
    63  
    64  func testPageBufferPoolWriteString(t *testing.T, pool parquet.PageBufferPool) {
    65  	const content = "Hello World!"
    66  
    67  	buffer := pool.GetPageBuffer()
    68  	_, err := io.WriteString(buffer, content)
    69  	if err != nil {
    70  		t.Fatal(err)
    71  	}
    72  
    73  	assertBufferContent(t, buffer, content)
    74  }
    75  
    76  func testPageBufferPoolCopyToBuffer(t *testing.T, pool parquet.PageBufferPool) {
    77  	const content = "ABC"
    78  
    79  	buffer := pool.GetPageBuffer()
    80  	reader := strings.NewReader(content)
    81  	_, err := io.Copy(buffer, struct{ io.Reader }{reader})
    82  	if err != nil {
    83  		t.Fatal(err)
    84  	}
    85  
    86  	assertBufferContent(t, buffer, content)
    87  }
    88  
    89  func testPageBufferPoolCopyFromBuffer(t *testing.T, pool parquet.PageBufferPool) {
    90  	const content = "0123456789"
    91  
    92  	buffer := pool.GetPageBuffer()
    93  	if _, err := io.WriteString(buffer, content); err != nil {
    94  		t.Fatal(err)
    95  	}
    96  
    97  	writer := new(bytes.Buffer)
    98  	_, err := io.Copy(struct{ io.Writer }{writer}, buffer)
    99  	if err != nil {
   100  		t.Fatal(err)
   101  	}
   102  
   103  	assertBufferContent(t, writer, content)
   104  }
   105  
   106  func assertBufferContent(t *testing.T, b io.Reader, s string) {
   107  	t.Helper()
   108  
   109  	if err := iotest.TestReader(b, []byte(s)); err != nil {
   110  		t.Error(err)
   111  	}
   112  }