github.com/grailbio/bigslice@v0.0.0-20230519005545-30c4c12152ad/exec/store_test.go (about)

     1  // Copyright 2018 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 exec
     6  
     7  import (
     8  	"bytes"
     9  	"context"
    10  	"io"
    11  	"io/ioutil"
    12  	"testing"
    13  
    14  	fuzz "github.com/google/gofuzz"
    15  	"github.com/grailbio/base/errors"
    16  	"github.com/grailbio/testutil"
    17  )
    18  
    19  func testStore(t *testing.T, store Store) {
    20  	t.Helper()
    21  	fz := fuzz.New()
    22  	fz.NumElements(1e3, 1e6)
    23  	var data []byte
    24  	fz.Fuzz(&data)
    25  	ctx := context.Background()
    26  	task := TaskName{Op: "test", Shard: 1, NumShard: 2}
    27  	wc, err := store.Create(ctx, task, 0)
    28  	if err != nil {
    29  		t.Error(err)
    30  		return
    31  	}
    32  	if _, err = io.Copy(wc, bytes.NewReader(data)); err != nil {
    33  		t.Error(err)
    34  		return
    35  	}
    36  	// Make sure the buffer isn't available until it's closed.
    37  	_, err = store.Open(ctx, task, 0, 0)
    38  	if err == nil {
    39  		t.Error("store prematurely unavailable")
    40  	} else if !errors.Is(errors.NotExist, err) {
    41  		t.Errorf("unexpected error: %v", err)
    42  	}
    43  	if err = wc.Commit(ctx, 12345); err != nil {
    44  		t.Error(err)
    45  		return
    46  	}
    47  	info, err := store.Stat(ctx, task, 0)
    48  	if err != nil {
    49  		t.Error(err)
    50  	} else {
    51  		if got, want := info.Size, int64(len(data)); got != want {
    52  			t.Errorf("got %v, want %v", got, want)
    53  		}
    54  		if got, want := info.Records, int64(12345); got != want {
    55  			t.Errorf("got %v, want %v", got, want)
    56  		}
    57  	}
    58  
    59  	rc, err := store.Open(ctx, task, 0, 0)
    60  	if err != nil {
    61  		t.Error(err)
    62  		return
    63  	}
    64  	defer rc.Close()
    65  	got, err := ioutil.ReadAll(rc)
    66  	if err != nil {
    67  		t.Error(err)
    68  		return
    69  	}
    70  	if !bytes.Equal(data, got) {
    71  		t.Error("data do not match")
    72  	}
    73  	// No partition 1 was ever stored, so discarding is a no-op.
    74  	err = store.Discard(ctx, task, 1)
    75  	if err == nil {
    76  		t.Error("expected error discarding non-existent (task, partition)")
    77  	}
    78  	// Make sure we can still Open successfully after the unrelated Discard.
    79  	rc, err = store.Open(ctx, task, 0, 0)
    80  	if err != nil {
    81  		t.Fatal(err)
    82  	}
    83  	rc.Close()
    84  	// Now discard, and try to Open. It should fail.
    85  	err = store.Discard(ctx, task, 0)
    86  	if err != nil {
    87  		t.Fatal(err)
    88  	}
    89  	_, err = store.Open(ctx, task, 0, 0)
    90  	if err == nil {
    91  		t.Fatal("expected error opening discarded (task, partition)")
    92  	}
    93  }
    94  
    95  func TestStoreImpls(t *testing.T) {
    96  	testStore(t, newMemoryStore())
    97  	dir, cleanup := testutil.TempDir(t, "", "")
    98  	defer cleanup()
    99  	testStore(t, &fileStore{dir})
   100  }