github.com/grailbio/base@v0.0.11/stateio/stateio_test.go (about)

     1  // Copyright 2019 GRAIL, Inc. All rights reserved.
     2  // Use of this source code is governed by the Apache 2.0
     3  // license that can be found in the LICENSE file.
     4  
     5  package stateio
     6  
     7  import (
     8  	"io"
     9  	"io/ioutil"
    10  	"math/rand"
    11  	"os"
    12  	"testing"
    13  )
    14  
    15  func TestStateIO(t *testing.T) {
    16  	file, cleanup := tempfile(t)
    17  	defer cleanup()
    18  
    19  	w := NewWriter(file, 0, 0)
    20  	must(t, w.Update(entry(1)))
    21  	must(t, w.Update(entry(2)))
    22  
    23  	state, _, updates, err := RestoreFile(file)
    24  	must(t, err)
    25  	if state != nil {
    26  		t.Fatal("non-nil snapshot")
    27  	}
    28  	b, err := updates.Read()
    29  	must(t, err)
    30  	mustEntry(t, 1, b)
    31  	b, err = updates.Read()
    32  	must(t, err)
    33  	mustEntry(t, 2, b)
    34  	mustEOF(t, updates)
    35  
    36  	must(t, file.Truncate(0))
    37  	state, _, updates, err = Restore(file, 0)
    38  	if state != nil {
    39  		t.Fatal("non-nil snapshot")
    40  	}
    41  	must(t, err)
    42  	mustEOF(t, updates)
    43  
    44  	const N = 100
    45  	w, err = NewFileWriter(file)
    46  	must(t, err)
    47  	for i := 0; i < N; i++ {
    48  		if i%10 == 0 {
    49  			must(t, w.Snapshot(entry(i)))
    50  		} else {
    51  			must(t, w.Update(entry(i)))
    52  		}
    53  		if i%5 == 0 {
    54  			// Reset the writer to make sure that writers
    55  			// resume properly.
    56  			must(t, err)
    57  			w, err = NewFileWriter(file)
    58  			must(t, err)
    59  		}
    60  	}
    61  
    62  	state, _, updates, err = RestoreFile(file)
    63  	must(t, err)
    64  	mustEntry(t, N-10, state)
    65  	for i := N - 9; i < N; i++ {
    66  		entry, err := updates.Read()
    67  		must(t, err)
    68  		mustEntry(t, i, entry)
    69  	}
    70  	mustEOF(t, updates)
    71  }
    72  
    73  func entry(n int) []byte {
    74  	b := make([]byte, n)
    75  	r := rand.New(rand.NewSource(int64(n)))
    76  	for i := range b {
    77  		b[i] = byte(r.Intn(256))
    78  	}
    79  	return b
    80  }
    81  
    82  func mustEntry(t *testing.T, n int, b []byte) {
    83  	t.Helper()
    84  	if got, want := len(b), n; got != want {
    85  		t.Fatalf("got %v, want %v", got, want)
    86  	}
    87  	r := rand.New(rand.NewSource(int64(n)))
    88  	for i := range b {
    89  		if got, want := int(b[i]), r.Intn(256); got != want {
    90  			t.Fatalf("byte %d: got %v, want %v", i, got, want)
    91  		}
    92  	}
    93  }
    94  
    95  func must(t *testing.T, err error) {
    96  	t.Helper()
    97  	if err != nil {
    98  		t.Fatal(err)
    99  	}
   100  }
   101  
   102  func mustEOF(t *testing.T, r *Reader) {
   103  	t.Helper()
   104  	_, err := r.Read()
   105  	if got, want := err, io.EOF; got != want {
   106  		t.Fatalf("got %v, want %v", got, want)
   107  	}
   108  }
   109  
   110  func tempfile(t *testing.T) (file *os.File, cleanup func()) {
   111  	t.Helper()
   112  	var err error
   113  	file, err = ioutil.TempFile("", "")
   114  	if err != nil {
   115  		t.Fatal(err)
   116  	}
   117  	os.Remove(file.Name())
   118  	cleanup = func() { file.Close() }
   119  	return
   120  }