github.com/treeverse/lakefs@v1.24.1-0.20240520134607-95648127bfb0/pkg/graveler/ref/branch_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/stretchr/testify/require"
    10  	"github.com/treeverse/lakefs/pkg/graveler"
    11  	"github.com/treeverse/lakefs/pkg/graveler/ref"
    12  	"github.com/treeverse/lakefs/pkg/testutil"
    13  )
    14  
    15  func TestBranchSimpleIterator(t *testing.T) {
    16  	r, kvStore := testRefManager(t)
    17  	branches := []graveler.BranchID{"a", "aa", "b", "c", "e", "d"}
    18  	ctx := context.Background()
    19  	repo := &graveler.RepositoryRecord{
    20  		RepositoryID: "repo1",
    21  		Repository: &graveler.Repository{
    22  			StorageNamespace: "s3://foo",
    23  			CreationDate:     time.Now(),
    24  			DefaultBranchID:  "main",
    25  		},
    26  	}
    27  	repository, err := r.CreateRepository(ctx, repo.RepositoryID, *repo.Repository)
    28  	testutil.Must(t, err)
    29  
    30  	// prepare data
    31  	for _, b := range branches {
    32  		testutil.Must(t, r.SetBranch(ctx, repository, b, graveler.Branch{CommitID: "c1"}))
    33  	}
    34  
    35  	t.Run("listing all branches", func(t *testing.T) {
    36  		iter, err := ref.NewBranchSimpleIterator(ctx, kvStore, repo)
    37  		require.NoError(t, err)
    38  		ids := make([]graveler.BranchID, 0)
    39  		for iter.Next() {
    40  			b := iter.Value()
    41  			ids = append(ids, b.BranchID)
    42  		}
    43  		if iter.Err() != nil {
    44  			t.Fatalf("unexpected error: %v", iter.Err())
    45  		}
    46  		iter.Close()
    47  
    48  		if diffs := deep.Equal(ids, []graveler.BranchID{"a", "aa", "b", "c", "d", "e", "main"}); diffs != nil {
    49  			t.Fatalf("got wrong list of IDs: %v", diffs)
    50  		}
    51  	})
    52  
    53  	t.Run("listing branches SeekGE", func(t *testing.T) {
    54  		iter, err := ref.NewBranchSimpleIterator(ctx, kvStore, repo)
    55  		require.NoError(t, err)
    56  		iter.SeekGE("b")
    57  		ids := make([]graveler.BranchID, 0)
    58  		for iter.Next() {
    59  			b := iter.Value()
    60  			ids = append(ids, b.BranchID)
    61  		}
    62  		if iter.Err() != nil {
    63  			t.Fatalf("unexpected error: %v", iter.Err())
    64  		}
    65  
    66  		if diffs := deep.Equal(ids, []graveler.BranchID{"b", "c", "d", "e", "main"}); diffs != nil {
    67  			t.Fatalf("got wrong list of branch IDs: %v", diffs)
    68  		}
    69  
    70  		// now let's seek
    71  		iter.SeekGE("aa")
    72  		ids = make([]graveler.BranchID, 0)
    73  		for iter.Next() {
    74  			b := iter.Value()
    75  			ids = append(ids, b.BranchID)
    76  		}
    77  		if iter.Err() != nil {
    78  			t.Fatalf("unexpected error: %v", iter.Err())
    79  		}
    80  		iter.Close()
    81  
    82  		if diffs := deep.Equal(ids, []graveler.BranchID{"aa", "b", "c", "d", "e", "main"}); diffs != nil {
    83  			t.Fatalf("got wrong list of branch IDs")
    84  		}
    85  
    86  		require.False(t, iter.Next())
    87  		require.Nil(t, iter.Value())
    88  		require.ErrorIs(t, iter.Err(), ref.ErrIteratorClosed)
    89  	})
    90  }
    91  
    92  func TestBranchByCommitIterator(t *testing.T) {
    93  	r, kvStore := testRefManager(t)
    94  	branches := []graveler.BranchID{"a", "aa", "b", "c", "e", "d"}
    95  	ctx := context.Background()
    96  	repo := &graveler.RepositoryRecord{
    97  		RepositoryID: "repo1",
    98  		Repository: &graveler.Repository{
    99  			StorageNamespace: "s3://foo",
   100  			CreationDate:     time.Now(),
   101  			DefaultBranchID:  "main",
   102  		},
   103  	}
   104  	repository, err := r.CreateRepository(ctx, repo.RepositoryID, *repo.Repository)
   105  	testutil.Must(t, err)
   106  
   107  	// prepare data
   108  	for i, b := range branches {
   109  		testutil.Must(t, r.SetBranch(ctx, repository, b, graveler.Branch{CommitID: graveler.CommitID(branches[len(branches)-i-1])}))
   110  	}
   111  
   112  	t.Run("listing all branches", func(t *testing.T) {
   113  		iter, err := ref.NewBranchByCommitIterator(ctx, kvStore, repo)
   114  		require.NoError(t, err)
   115  		ids := []graveler.CommitID{"mainCommitNotFound"}
   116  		for iter.Next() {
   117  			b := iter.Value()
   118  			// save default branch commit ID since its created randomly as the first element of ids
   119  			if b.BranchID.String() == "main" {
   120  				ids[0] = b.CommitID
   121  			} else {
   122  				ids = append(ids, b.CommitID)
   123  			}
   124  		}
   125  		if iter.Err() != nil {
   126  			t.Fatalf("unexpected error: %v", iter.Err())
   127  		}
   128  		iter.Close()
   129  
   130  		if diffs := deep.Equal(ids, []graveler.CommitID{ids[0], "a", "aa", "b", "c", "d", "e"}); diffs != nil {
   131  			t.Fatalf("got wrong list of IDs: %v", diffs)
   132  		}
   133  	})
   134  }