github.com/pingcap/br@v5.3.0-alpha.0.20220125034240-ec59c7b6ce30+incompatible/pkg/storage/writer_test.go (about)

     1  // Copyright 2020 PingCAP, Inc. Licensed under Apache-2.0.
     2  
     3  package storage
     4  
     5  import (
     6  	"bytes"
     7  	"context"
     8  	"io"
     9  	"os"
    10  	"path/filepath"
    11  	"strings"
    12  
    13  	. "github.com/pingcap/check"
    14  )
    15  
    16  func (r *testStorageSuite) TestExternalFileWriter(c *C) {
    17  	dir := c.MkDir()
    18  
    19  	type testcase struct {
    20  		name    string
    21  		content []string
    22  	}
    23  	testFn := func(test *testcase, c *C) {
    24  		c.Log(test.name)
    25  		backend, err := ParseBackend("local://"+filepath.ToSlash(dir), nil)
    26  		c.Assert(err, IsNil)
    27  		ctx := context.Background()
    28  		storage, err := Create(ctx, backend, true)
    29  		c.Assert(err, IsNil)
    30  		fileName := strings.ReplaceAll(test.name, " ", "-") + ".txt"
    31  		writer, err := storage.Create(ctx, fileName)
    32  		c.Assert(err, IsNil)
    33  		for _, str := range test.content {
    34  			p := []byte(str)
    35  			written, err2 := writer.Write(ctx, p)
    36  			c.Assert(err2, IsNil)
    37  			c.Assert(written, Equals, len(p))
    38  		}
    39  		err = writer.Close(ctx)
    40  		c.Assert(err, IsNil)
    41  		content, err := os.ReadFile(filepath.Join(dir, fileName))
    42  		c.Assert(err, IsNil)
    43  		c.Assert(string(content), Equals, strings.Join(test.content, ""))
    44  	}
    45  	tests := []testcase{
    46  		{
    47  			name:    "short and sweet",
    48  			content: []string{"hi"},
    49  		},
    50  		{
    51  			name: "long text small chunks",
    52  			content: []string{
    53  				"hello world",
    54  				"hello world",
    55  				"hello world",
    56  				"hello world",
    57  				"hello world",
    58  				"hello world",
    59  			},
    60  		},
    61  		{
    62  			name: "long text medium chunks",
    63  			content: []string{
    64  				"hello world",
    65  				"hello world",
    66  				"hello world",
    67  				"hello world",
    68  				"hello world",
    69  				"hello world",
    70  			},
    71  		},
    72  		{
    73  			name: "long text large chunks",
    74  			content: []string{
    75  				"hello world",
    76  				"hello world",
    77  				"hello world",
    78  				"hello world",
    79  				"hello world",
    80  				"hello world",
    81  			},
    82  		},
    83  	}
    84  	for i := range tests {
    85  		testFn(&tests[i], c)
    86  	}
    87  }
    88  
    89  func (r *testStorageSuite) TestCompressReaderWriter(c *C) {
    90  	dir := c.MkDir()
    91  
    92  	type testcase struct {
    93  		name         string
    94  		content      []string
    95  		compressType CompressType
    96  	}
    97  	testFn := func(test *testcase, c *C) {
    98  		c.Log(test.name)
    99  		backend, err := ParseBackend("local://"+filepath.ToSlash(dir), nil)
   100  		c.Assert(err, IsNil)
   101  		ctx := context.Background()
   102  		storage, err := Create(ctx, backend, true)
   103  		c.Assert(err, IsNil)
   104  		storage = WithCompression(storage, Gzip)
   105  		fileName := strings.ReplaceAll(test.name, " ", "-") + ".txt.gz"
   106  		writer, err := storage.Create(ctx, fileName)
   107  		c.Assert(err, IsNil)
   108  		for _, str := range test.content {
   109  			p := []byte(str)
   110  			written, err2 := writer.Write(ctx, p)
   111  			c.Assert(err2, IsNil)
   112  			c.Assert(written, Equals, len(p))
   113  		}
   114  		err = writer.Close(ctx)
   115  		c.Assert(err, IsNil)
   116  
   117  		// make sure compressed file is written correctly
   118  		file, err := os.Open(filepath.Join(dir, fileName))
   119  		c.Assert(err, IsNil)
   120  		r, err := newCompressReader(test.compressType, file)
   121  		c.Assert(err, IsNil)
   122  		var bf bytes.Buffer
   123  		_, err = bf.ReadFrom(r)
   124  		c.Assert(err, IsNil)
   125  		c.Assert(bf.String(), Equals, strings.Join(test.content, ""))
   126  		c.Assert(r.Close(), IsNil)
   127  
   128  		// test withCompression Open
   129  		r, err = storage.Open(ctx, fileName)
   130  		c.Assert(err, IsNil)
   131  		content, err := io.ReadAll(r)
   132  		c.Assert(err, IsNil)
   133  		c.Assert(string(content), Equals, strings.Join(test.content, ""))
   134  
   135  		c.Assert(file.Close(), IsNil)
   136  	}
   137  	compressTypeArr := []CompressType{Gzip}
   138  	tests := []testcase{
   139  		{
   140  			name: "long text medium chunks",
   141  			content: []string{
   142  				"hello world",
   143  				"hello world",
   144  				"hello world",
   145  				"hello world",
   146  				"hello world",
   147  				"hello world",
   148  			},
   149  		},
   150  		{
   151  			name: "long text large chunks",
   152  			content: []string{
   153  				"hello world",
   154  				"hello world",
   155  				"hello world",
   156  				"hello world",
   157  				"hello world",
   158  				"hello world",
   159  			},
   160  		},
   161  	}
   162  	for i := range tests {
   163  		for _, compressType := range compressTypeArr {
   164  			tests[i].compressType = compressType
   165  			testFn(&tests[i], c)
   166  		}
   167  	}
   168  }