github.com/treeverse/lakefs@v1.24.1-0.20240520134607-95648127bfb0/pkg/graveler/ref/repository_iterator_test.go (about)

     1  package ref_test
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/go-test/deep"
     9  	"github.com/golang/mock/gomock"
    10  	"github.com/stretchr/testify/require"
    11  	"github.com/treeverse/lakefs/pkg/graveler"
    12  	"github.com/treeverse/lakefs/pkg/graveler/ref"
    13  	"github.com/treeverse/lakefs/pkg/kv/mock"
    14  	"github.com/treeverse/lakefs/pkg/testutil"
    15  )
    16  
    17  func TestRepositoryIterator(t *testing.T) {
    18  	r, store := testRefManager(t)
    19  	repos := []graveler.RepositoryID{"a", "aa", "b", "c", "e", "d", "f"}
    20  
    21  	// prepare data
    22  	for _, repoId := range repos {
    23  		_, err := r.CreateRepository(context.Background(), repoId, graveler.Repository{
    24  			StorageNamespace: "s3://foo",
    25  			CreationDate:     time.Now(),
    26  			DefaultBranchID:  "main",
    27  		})
    28  		testutil.Must(t, err)
    29  	}
    30  
    31  	t.Run("listing all repos", func(t *testing.T) {
    32  		iter, err := ref.NewRepositoryIterator(context.Background(), store)
    33  		require.NoError(t, err)
    34  		repoIds := make([]graveler.RepositoryID, 0)
    35  		for iter.Next() {
    36  			repo := iter.Value()
    37  			repoIds = append(repoIds, repo.RepositoryID)
    38  		}
    39  		if iter.Err() != nil {
    40  			t.Fatalf("unexpected error: %v", iter.Err())
    41  		}
    42  		iter.Close()
    43  
    44  		if diffs := deep.Equal(repoIds, []graveler.RepositoryID{"a", "aa", "b", "c", "d", "e", "f"}); diffs != nil {
    45  			t.Fatalf("got wrong list of repo IDs: %v", diffs)
    46  		}
    47  	})
    48  
    49  	t.Run("listing repos from prefix", func(t *testing.T) {
    50  		iter, err := ref.NewRepositoryIterator(context.Background(), store)
    51  		require.NoError(t, err)
    52  		iter.SeekGE("b")
    53  		repoIds := make([]graveler.RepositoryID, 0)
    54  		for iter.Next() {
    55  			repo := iter.Value()
    56  			repoIds = append(repoIds, repo.RepositoryID)
    57  		}
    58  		if iter.Err() != nil {
    59  			t.Fatalf("unexpected error: %v", iter.Err())
    60  		}
    61  		iter.Close()
    62  
    63  		if diffs := deep.Equal(repoIds, []graveler.RepositoryID{"b", "c", "d", "e", "f"}); diffs != nil {
    64  			t.Fatalf("got wrong list of repo IDs: %v", diffs)
    65  		}
    66  	})
    67  
    68  	t.Run("listing repos SeekGE", func(t *testing.T) {
    69  		iter, err := ref.NewRepositoryIterator(context.Background(), store)
    70  		require.NoError(t, err)
    71  		iter.SeekGE("b")
    72  		repoIds := make([]graveler.RepositoryID, 0)
    73  		for iter.Next() {
    74  			repo := iter.Value()
    75  			repoIds = append(repoIds, repo.RepositoryID)
    76  		}
    77  		if iter.Err() != nil {
    78  			t.Fatalf("unexpected error: %v", iter.Err())
    79  		}
    80  
    81  		if diffs := deep.Equal(repoIds, []graveler.RepositoryID{"b", "c", "d", "e", "f"}); diffs != nil {
    82  			t.Fatalf("got wrong list of repo IDs: %v", diffs)
    83  		}
    84  
    85  		// now let's seek
    86  		iter.SeekGE("aa")
    87  
    88  		repoIds = make([]graveler.RepositoryID, 0)
    89  		for iter.Next() {
    90  			repo := iter.Value()
    91  			repoIds = append(repoIds, repo.RepositoryID)
    92  		}
    93  		if iter.Err() != nil {
    94  			t.Fatalf("unexpected error: %v", iter.Err())
    95  		}
    96  		iter.Close()
    97  
    98  		if diffs := deep.Equal(repoIds, []graveler.RepositoryID{"aa", "b", "c", "d", "e", "f"}); diffs != nil {
    99  			t.Fatalf("got wrong list of repo IDs: %v", diffs)
   100  		}
   101  	})
   102  }
   103  
   104  func TestRepositoryIterator_CloseTwice(t *testing.T) {
   105  	ctx := context.Background()
   106  	ctrl := gomock.NewController(t)
   107  	entIt := mock.NewMockEntriesIterator(ctrl)
   108  	entIt.EXPECT().Close().Times(1)
   109  	store := mock.NewMockStore(ctrl)
   110  	store.EXPECT().Scan(ctx, gomock.Any(), gomock.Any()).Return(entIt, nil).Times(1)
   111  	it, err := ref.NewRepositoryIterator(ctx, store)
   112  	if err != nil {
   113  		t.Fatal("NewRepositoryIterator failed", err)
   114  	}
   115  	it.Close()
   116  	// Make sure calling Close again do not crash
   117  	it.Close()
   118  }
   119  
   120  func TestRepositoryIterator_NextClosed(t *testing.T) {
   121  	ctx := context.Background()
   122  	ctrl := gomock.NewController(t)
   123  	entIt := mock.NewMockEntriesIterator(ctrl)
   124  	entIt.EXPECT().Close().Times(1)
   125  	store := mock.NewMockStore(ctrl)
   126  	store.EXPECT().Scan(ctx, gomock.Any(), gomock.Any()).Return(entIt, nil).Times(1)
   127  	it, err := ref.NewRepositoryIterator(ctx, store)
   128  	if err != nil {
   129  		t.Fatal("NewRepositoryIterator failed", err)
   130  	}
   131  	it.Close()
   132  	// Make sure calling Next should not crash
   133  	it.Next()
   134  }