1 files changed, 227 insertions, 0 deletions
diff --git a/pkg/compressio/compressio_test.go b/pkg/compressio/compressio_test.go
new file mode 100644
index 000000000..d7911419d
--- /dev/null
+++ b/pkg/compressio/compressio_test.go
@@ -0,0 +1,227 @@
+// Copyright 2018 Google Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package compressio
+
+import (
+	"bytes"
+	"compress/flate"
+	"encoding/base64"
+	"fmt"
+	"io"
+	"math/rand"
+	"runtime"
+	"testing"
+	"time"
+)
+
+type harness interface {
+	Errorf(format string, v ...interface{})
+	Fatalf(format string, v ...interface{})
+	Logf(format string, v ...interface{})
+}
+
+func initTest(t harness, size int) []byte {
+	// Set number of processes to number of CPUs.
+	runtime.GOMAXPROCS(runtime.NumCPU())
+
+	// Construct synthetic data. We do this by encoding random data with
+	// base64. This gives a high level of entropy, but still quite a bit of
+	// structure, to give reasonable compression ratios (~75%).
+	var buf bytes.Buffer
+	bufW := base64.NewEncoder(base64.RawStdEncoding, &buf)
+	bufR := rand.New(rand.NewSource(0))
+	if _, err := io.CopyN(bufW, bufR, int64(size)); err != nil {
+		t.Fatalf("unable to seed random data: %v", err)
+	}
+	return buf.Bytes()
+}
+
+type testOpts struct {
+	Name            string
+	Data            []byte
+	NewWriter       func(*bytes.Buffer) (io.Writer, error)
+	NewReader       func(*bytes.Buffer) (io.Reader, error)
+	PreCompress     func()
+	PostCompress    func()
+	PreDecompress   func()
+	PostDecompress  func()
+	CompressIters   int
+	DecompressIters int
+}
+
+func doTest(t harness, opts testOpts) {
+	// Compress.
+	var compressed bytes.Buffer
+	compressionStartTime := time.Now()
+	if opts.PreCompress != nil {
+		opts.PreCompress()
+	}
+	if opts.CompressIters <= 0 {
+		opts.CompressIters = 1
+	}
+	for i := 0; i < opts.CompressIters; i++ {
+		compressed.Reset()
+		w, err := opts.NewWriter(&compressed)
+		if err != nil {
+			t.Errorf("%s: NewWriter got err %v, expected nil", opts.Name, err)
+		}
+		if _, err := io.Copy(w, bytes.NewBuffer(opts.Data)); err != nil {
+			t.Errorf("%s: compress got err %v, expected nil", opts.Name, err)
+			return
+		}
+		closer, ok := w.(io.Closer)
+		if ok {
+			if err := closer.Close(); err != nil {
+				t.Errorf("%s: got err %v, expected nil", opts.Name, err)
+				return
+			}
+		}
+	}
+	if opts.PostCompress != nil {
+		opts.PostCompress()
+	}
+	compressionTime := time.Since(compressionStartTime)
+	compressionRatio := float32(compressed.Len()) / float32(len(opts.Data))
+
+	// Decompress.
+	var decompressed bytes.Buffer
+	decompressionStartTime := time.Now()
+	if opts.PreDecompress != nil {
+		opts.PreDecompress()
+	}
+	if opts.DecompressIters <= 0 {
+		opts.DecompressIters = 1
+	}
+	for i := 0; i < opts.DecompressIters; i++ {
+		decompressed.Reset()
+		r, err := opts.NewReader(bytes.NewBuffer(compressed.Bytes()))
+		if err != nil {
+			t.Errorf("%s: NewReader got err %v, expected nil", opts.Name, err)
+			return
+		}
+		if _, err := io.Copy(&decompressed, r); err != nil {
+			t.Errorf("%s: decompress got err %v, expected nil", opts.Name, err)
+			return
+		}
+	}
+	if opts.PostDecompress != nil {
+		opts.PostDecompress()
+	}
+	decompressionTime := time.Since(decompressionStartTime)
+
+	// Verify.
+	if decompressed.Len() != len(opts.Data) {
+		t.Errorf("%s: got %d bytes, expected %d", opts.Name, decompressed.Len(), len(opts.Data))
+	}
+	if !bytes.Equal(opts.Data, decompressed.Bytes()) {
+		t.Errorf("%s: got mismatch, expected match", opts.Name)
+		if len(opts.Data) < 32 { // Don't flood the logs.
+			t.Errorf("got %v, expected %v", decompressed.Bytes(), opts.Data)
+		}
+	}
+
+	t.Logf("%s: compression time %v, ratio %2.2f, decompression time %v",
+		opts.Name, compressionTime, compressionRatio, decompressionTime)
+}
+
+func TestCompress(t *testing.T) {
+	var (
+		data  = initTest(t, 10*1024*1024)
+		data0 = data[:0]
+		data1 = data[:1]
+		data2 = data[:11]
+		data3 = data[:16]
+		data4 = data[:]
+	)
+
+	for _, data := range [][]byte{data0, data1, data2, data3, data4} {
+		for _, blockSize := range []uint32{1, 4, 1024, 4 * 1024, 16 * 1024} {
+			// Skip annoying tests; they just take too long.
+			if blockSize <= 16 && len(data) > 16 {
+				continue
+			}
+
+			// Do the compress test.
+			doTest(t, testOpts{
+				Name: fmt.Sprintf("len(data)=%d, blockSize=%d", len(data), blockSize),
+				Data: data,
+				NewWriter: func(b *bytes.Buffer) (io.Writer, error) {
+					return NewWriter(b, blockSize, flate.BestCompression)
+				},
+				NewReader: func(b *bytes.Buffer) (io.Reader, error) {
+					return NewReader(b)
+				},
+			})
+		}
+
+		// Do the vanilla test.
+		doTest(t, testOpts{
+			Name: fmt.Sprintf("len(data)=%d, vanilla flate", len(data)),
+			Data: data,
+			NewWriter: func(b *bytes.Buffer) (io.Writer, error) {
+				return flate.NewWriter(b, flate.BestCompression)
+			},
+			NewReader: func(b *bytes.Buffer) (io.Reader, error) {
+				return flate.NewReader(b), nil
+			},
+		})
+	}
+}
+
+const (
+	// benchBlockSize is the blockSize for benchmarks.
+	benchBlockSize = 32 * 1024
+
+	// benchDataSize is the amount of data for benchmarks.
+	benchDataSize = 10 * 1024 * 1024
+)
+
+func BenchmarkCompress(b *testing.B) {
+	b.StopTimer()
+	b.SetBytes(benchDataSize)
+	data := initTest(b, benchDataSize)
+	doTest(b, testOpts{
+		Name:         fmt.Sprintf("len(data)=%d, blockSize=%d", len(data), benchBlockSize),
+		Data:         data,
+		PreCompress:  b.StartTimer,
+		PostCompress: b.StopTimer,
+		NewWriter: func(b *bytes.Buffer) (io.Writer, error) {
+			return NewWriter(b, benchBlockSize, flate.BestCompression)
+		},
+		NewReader: func(b *bytes.Buffer) (io.Reader, error) {
+			return NewReader(b)
+		},
+		CompressIters: b.N,
+	})
+}
+
+func BenchmarkDecompress(b *testing.B) {
+	b.StopTimer()
+	b.SetBytes(benchDataSize)
+	data := initTest(b, benchDataSize)
+	doTest(b, testOpts{
+		Name:           fmt.Sprintf("len(data)=%d, blockSize=%d", len(data), benchBlockSize),
+		Data:           data,
+		PreDecompress:  b.StartTimer,
+		PostDecompress: b.StopTimer,
+		NewWriter: func(b *bytes.Buffer) (io.Writer, error) {
+			return NewWriter(b, benchBlockSize, flate.BestCompression)
+		},
+		NewReader: func(b *bytes.Buffer) (io.Reader, error) {
+			return NewReader(b)
+		},
+		DecompressIters: b.N,
+	})
+}