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

     1  package kv_test
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"testing"
     7  
     8  	"github.com/golang/mock/gomock"
     9  	"github.com/treeverse/lakefs/pkg/kv"
    10  	"github.com/treeverse/lakefs/pkg/kv/kvtest"
    11  	"github.com/treeverse/lakefs/pkg/kv/mock"
    12  )
    13  
    14  func TestPartitionIterator_ClosedBehaviour(t *testing.T) {
    15  	ctx := context.Background()
    16  	ctrl := gomock.NewController(t)
    17  	store := mock.NewMockStore(ctrl)
    18  	entIt := mock.NewMockEntriesIterator(ctrl)
    19  	entIt.EXPECT().Close().Times(1)
    20  	store.EXPECT().Scan(ctx, gomock.Any(), gomock.Any()).Return(entIt, nil).Times(1)
    21  
    22  	it := kv.NewPartitionIterator(ctx, store, (&kvtest.TestModel{}).ProtoReflect().Type(), "partitionKey", 0)
    23  	it.SeekGE([]byte("test"))
    24  	it.Close()
    25  
    26  	err := it.Err() // verify we don't crash or call underlying iterator after we Close()
    27  	if err != nil {
    28  		t.Fatal("Err() on closed iterator", err)
    29  	}
    30  }
    31  
    32  func TestPartitionIterator_CloseAfterSeekGEFailed(t *testing.T) {
    33  	ctx := context.Background()
    34  	ctrl := gomock.NewController(t)
    35  	store := mock.NewMockStore(ctrl)
    36  	entIt := mock.NewMockEntriesIterator(ctrl)
    37  	entItErr := errors.New("failed to scan")
    38  	store.EXPECT().Scan(ctx, gomock.Any(), gomock.Any()).Return(entIt, entItErr).Times(1)
    39  	entIt.EXPECT().Close().Times(1)
    40  	it := kv.NewPartitionIterator(ctx, store, (&kvtest.TestModel{}).ProtoReflect().Type(), "partitionKey", 0)
    41  	it.SeekGE([]byte("test"))
    42  	it.Close() // verify we don't crash after SeekGE failed internally with Scan
    43  }