github.com/weaviate/weaviate@v1.24.6/adapters/repos/db/roaringset/segment_cursor_test.go (about)

     1  //                           _       _
     2  // __      _____  __ ___   ___  __ _| |_ ___
     3  // \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
     4  //  \ V  V /  __/ (_| |\ V /| | (_| | ||  __/
     5  //   \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
     6  //
     7  //  Copyright © 2016 - 2024 Weaviate B.V. All rights reserved.
     8  //
     9  //  CONTACT: hello@weaviate.io
    10  //
    11  
    12  package roaringset
    13  
    14  import (
    15  	"fmt"
    16  	"testing"
    17  
    18  	"github.com/stretchr/testify/assert"
    19  	"github.com/stretchr/testify/require"
    20  	"github.com/weaviate/weaviate/adapters/repos/db/lsmkv/segmentindex"
    21  )
    22  
    23  func TestSegmentCursor(t *testing.T) {
    24  	seg, offsets := createDummySegment(t, 5)
    25  
    26  	t.Run("starting from beginning", func(t *testing.T) {
    27  		c := NewSegmentCursor(seg, nil)
    28  		key, layer, err := c.First()
    29  		require.Nil(t, err)
    30  		assert.Equal(t, []byte("00000"), key)
    31  		assert.True(t, layer.Additions.Contains(0))
    32  		assert.True(t, layer.Additions.Contains(1))
    33  		assert.True(t, layer.Deletions.Contains(2))
    34  		assert.True(t, layer.Deletions.Contains(3))
    35  	})
    36  
    37  	t.Run("starting from beginning, page through all", func(t *testing.T) {
    38  		c := NewSegmentCursor(seg, nil)
    39  		it := uint64(0)
    40  		for key, layer, err := c.First(); key != nil; key, layer, err = c.Next() {
    41  			require.Nil(t, err)
    42  			assert.Equal(t, []byte(fmt.Sprintf("%05d", it)), key)
    43  			assert.True(t, layer.Additions.Contains(it*4))
    44  			assert.True(t, layer.Additions.Contains(it*4+1))
    45  			assert.True(t, layer.Deletions.Contains(it*4+2))
    46  			assert.True(t, layer.Deletions.Contains(it*4+3))
    47  			it++
    48  		}
    49  
    50  		assert.Equal(t, uint64(5), it)
    51  	})
    52  
    53  	t.Run("seek and iterate from there", func(t *testing.T) {
    54  		seeker := createDummySeeker(t, offsets, 3)
    55  		c := NewSegmentCursor(seg, seeker)
    56  
    57  		// start on it 3 as this is where the seeker points us
    58  		it := uint64(3)
    59  		for key, layer, err := c.Seek([]byte("dummyseeker")); key != nil; key, layer, err = c.Next() {
    60  			require.Nil(t, err)
    61  			assert.Equal(t, []byte(fmt.Sprintf("%05d", it)), key)
    62  			assert.True(t, layer.Additions.Contains(it*4))
    63  			assert.True(t, layer.Additions.Contains(it*4+1))
    64  			assert.True(t, layer.Deletions.Contains(it*4+2))
    65  			assert.True(t, layer.Deletions.Contains(it*4+3))
    66  			it++
    67  		}
    68  
    69  		assert.Equal(t, uint64(5), it)
    70  	})
    71  
    72  	t.Run("seeker returns error", func(t *testing.T) {
    73  		seeker := createDummySeeker(t, offsets, 3)
    74  		seeker.err = fmt.Errorf("seek and fail")
    75  		c := NewSegmentCursor(seg, seeker)
    76  
    77  		_, _, err := c.Seek([]byte("dummyseeker"))
    78  		require.NotNil(t, err)
    79  		assert.Contains(t, err.Error(), "seek and fail")
    80  	})
    81  }
    82  
    83  func createDummySegment(t *testing.T, count uint64) ([]byte, []uint64) {
    84  	out := []byte{}
    85  	offsets := []uint64{}
    86  
    87  	for i := uint64(0); i < count; i++ {
    88  		key := []byte(fmt.Sprintf("%05d", i))
    89  		add := NewBitmap(i*4, i*4+1)
    90  		del := NewBitmap(i*4+2, i*4+3)
    91  		sn, err := NewSegmentNode(key, add, del)
    92  		require.Nil(t, err)
    93  		offsets = append(offsets, uint64(len(out)))
    94  		out = append(out, sn.ToBuffer()...)
    95  	}
    96  
    97  	return out, offsets
    98  }
    99  
   100  func createDummySeeker(t *testing.T, offsets []uint64, pos int) *dummySeeker {
   101  	return &dummySeeker{offsets, pos, nil}
   102  }
   103  
   104  type dummySeeker struct {
   105  	offsets []uint64
   106  	pos     int
   107  	err     error
   108  }
   109  
   110  // Seek returns the hard-coded pos that was set on init time, it ignores the
   111  // key
   112  func (s dummySeeker) Seek(key []byte) (segmentindex.Node, error) {
   113  	return segmentindex.Node{
   114  		Start: s.offsets[s.pos],
   115  		End:   s.offsets[s.pos+1],
   116  	}, s.err
   117  }