github.com/treeverse/lakefs@v1.24.1-0.20240520134607-95648127bfb0/pkg/kv/msg_test.go (about)

     1  package kv_test
     2  
     3  import (
     4  	"context"
     5  	"log"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/ory/dockertest/v3"
    10  	"github.com/stretchr/testify/require"
    11  	"github.com/treeverse/lakefs/pkg/kv"
    12  	"github.com/treeverse/lakefs/pkg/kv/kvparams"
    13  	"github.com/treeverse/lakefs/pkg/kv/kvtest"
    14  	_ "github.com/treeverse/lakefs/pkg/kv/mem"
    15  	"github.com/treeverse/lakefs/pkg/kv/postgres"
    16  	"github.com/treeverse/lakefs/pkg/testutil"
    17  	"google.golang.org/protobuf/proto"
    18  	"google.golang.org/protobuf/types/known/timestamppb"
    19  )
    20  
    21  const modelPartitionKey = "tm"
    22  
    23  func TestMsgFuncs(t *testing.T) {
    24  	ctx := context.Background()
    25  	store := kvtest.GetStore(ctx, t)
    26  
    27  	t.Run("set-get", func(t *testing.T) {
    28  		testSetGetMsg(t, ctx, store)
    29  	})
    30  	t.Run("set-if", func(t *testing.T) {
    31  		testSetIfMsg(t, ctx, store)
    32  	})
    33  }
    34  
    35  func testSetGetMsg(t testing.TB, ctx context.Context, store kv.Store) {
    36  	// set model info
    37  	setModel := &kvtest.TestModel{
    38  		Name:        []byte("SetGetModel"),
    39  		JustAString: "This is another string",
    40  		ADouble:     2.4,
    41  		TestTime:    timestamppb.New(time.Now().UTC()),
    42  		TestMap: map[string]int32{
    43  			"one":   1,
    44  			"two":   2,
    45  			"three": 3,
    46  		},
    47  		TestList: []bool{true, true, false, true, false},
    48  	}
    49  	err := kv.SetMsg(ctx, store, modelPartitionKey, setModel.Name, setModel)
    50  	if err != nil {
    51  		t.Fatal("failed to set model", err)
    52  	}
    53  
    54  	// get model info
    55  	m := &kvtest.TestModel{}
    56  	_, err = kv.GetMsg(ctx, store, modelPartitionKey, setModel.Name, m)
    57  	if err != nil {
    58  		t.Fatal("failed to get message", err)
    59  	}
    60  	equals := proto.Equal(m, setModel)
    61  	if !equals {
    62  		t.Errorf("Get model not equal: %s, expected: %s", m, setModel)
    63  	}
    64  }
    65  
    66  func testSetIfMsg(t testing.TB, ctx context.Context, store kv.Store) {
    67  	// set model info
    68  	setModel := &kvtest.TestModel{
    69  		Name:        []byte("SetIfModel"),
    70  		JustAString: "This is another string",
    71  		ADouble:     2.4,
    72  		TestTime:    timestamppb.New(time.Now()),
    73  		TestMap: map[string]int32{
    74  			"one":   1,
    75  			"two":   2,
    76  			"three": 3,
    77  		},
    78  		TestList: []bool{true, true, false, true, false},
    79  	}
    80  	err := kv.SetMsgIf(ctx, store, modelPartitionKey, setModel.Name, setModel, nil)
    81  	if err != nil {
    82  		t.Fatal("failed to set model with nil predicate", err)
    83  	}
    84  	pred, err := kv.GetMsg(ctx, store, modelPartitionKey, setModel.Name, nil)
    85  	if err != nil {
    86  		t.Fatal("failed to get model without msg", err)
    87  	}
    88  
    89  	// SetIf model
    90  	m1 := &kvtest.TestModel{
    91  		Name:        setModel.Name,
    92  		JustAString: "just another string",
    93  		ADouble:     3.14159,
    94  		TestTime:    timestamppb.New(time.Now().UTC()),
    95  		TestMap: map[string]int32{
    96  			"red":   1,
    97  			"green": 2,
    98  			"blue":  3,
    99  		},
   100  		TestList: []bool{true},
   101  	}
   102  
   103  	// SetMsgIf fails nil
   104  	err = kv.SetMsgIf(ctx, store, modelPartitionKey, setModel.Name, m1, nil)
   105  	require.Error(t, kv.ErrPredicateFailed, err)
   106  
   107  	// get model info
   108  	m2 := &kvtest.TestModel{}
   109  	_, err = kv.GetMsg(ctx, store, modelPartitionKey, setModel.Name, m2)
   110  	if err != nil {
   111  		t.Fatal("failed to get message", err)
   112  	}
   113  	equals := proto.Equal(m2, setModel)
   114  	if !equals {
   115  		t.Errorf("Get model not equal: %s, expected: %s", m2, setModel)
   116  	}
   117  
   118  	// SetIf succeeds
   119  	err = kv.SetMsgIf(ctx, store, modelPartitionKey, setModel.Name, m1, pred)
   120  	if err != nil {
   121  		t.Fatal("failed on SetIf", err)
   122  	}
   123  
   124  	_, err = kv.GetMsg(ctx, store, modelPartitionKey, setModel.Name, m2)
   125  	if err != nil {
   126  		t.Fatal("failed to get message", err)
   127  	}
   128  	equals = proto.Equal(m2, m1)
   129  	if !equals {
   130  		t.Errorf("Get model not equal: %s, expected: %s", m2, m1)
   131  	}
   132  
   133  	// Cleanup
   134  	testutil.MustDo(t, "cleanup", store.Delete(ctx, []byte(modelPartitionKey), setModel.Name))
   135  }
   136  
   137  func BenchmarkDrivers(b *testing.B) {
   138  	ctx := context.Background()
   139  
   140  	pool, err := dockertest.NewPool("")
   141  	if err != nil {
   142  		log.Fatalf("Could not connect to Docker: %s", err)
   143  	}
   144  
   145  	databaseURI, closer := testutil.GetDBInstance(pool)
   146  	defer closer()
   147  
   148  	dynamoStore := testutil.GetDynamoDBProd(ctx, b)
   149  	postgresStore, err := kv.Open(ctx, kvparams.Config{Type: postgres.DriverName, Postgres: &kvparams.Postgres{ConnectionString: databaseURI}})
   150  	if err != nil {
   151  		b.Fatal("failed to open postgres store", err)
   152  	}
   153  	defer postgresStore.Close()
   154  
   155  	tests := []struct {
   156  		name  string
   157  		store kv.Store
   158  	}{
   159  		{
   160  			name:  "postgres",
   161  			store: postgresStore,
   162  		},
   163  		{
   164  			name:  "dynamoDB",
   165  			store: dynamoStore,
   166  		},
   167  	}
   168  	for _, tt := range tests {
   169  		b.Run(tt.name, func(b *testing.B) {
   170  			b.ResetTimer()
   171  			for n := 0; n < b.N; n++ {
   172  				testSetGetMsg(b, ctx, tt.store)
   173  				testSetIfMsg(b, ctx, tt.store)
   174  			}
   175  		})
   176  	}
   177  }