github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/ingester/streams_map_test.go (about)

     1  package ingester
     2  
     3  import (
     4  	"testing"
     5  
     6  	"github.com/prometheus/common/model"
     7  	"github.com/prometheus/prometheus/model/labels"
     8  	"github.com/stretchr/testify/require"
     9  
    10  	"github.com/grafana/loki/pkg/validation"
    11  )
    12  
    13  func TestStreamsMap(t *testing.T) {
    14  	limits, err := validation.NewOverrides(defaultLimitsTestConfig(), nil)
    15  	require.NoError(t, err)
    16  	limiter := NewLimiter(limits, NilMetrics, &ringCountMock{count: 1}, 1)
    17  
    18  	ss := []*stream{
    19  		newStream(
    20  			defaultConfig(),
    21  			limiter,
    22  			"fake",
    23  			model.Fingerprint(1),
    24  			labels.Labels{
    25  				{Name: "foo", Value: "bar"},
    26  			},
    27  			true,
    28  			NilMetrics,
    29  		),
    30  		newStream(
    31  			defaultConfig(),
    32  			limiter,
    33  			"fake",
    34  			model.Fingerprint(2),
    35  			labels.Labels{
    36  				{Name: "bar", Value: "foo"},
    37  			},
    38  			true,
    39  			NilMetrics,
    40  		),
    41  	}
    42  	var s *stream
    43  	var loaded bool
    44  
    45  	streams := newStreamsMap()
    46  
    47  	require.Equal(t, 0, streams.Len())
    48  	s, loaded = streams.Load(ss[0].labelsString)
    49  	require.Nil(t, s)
    50  	require.False(t, loaded)
    51  	s, loaded = streams.LoadByFP(ss[1].fp)
    52  	require.Nil(t, s)
    53  	require.False(t, loaded)
    54  
    55  	// Test LoadOrStoreNew
    56  	s, loaded, err = streams.LoadOrStoreNew(ss[0].labelsString, func() (*stream, error) {
    57  		return ss[0], nil
    58  	}, nil)
    59  	require.Equal(t, s, ss[0])
    60  	require.False(t, loaded)
    61  	require.Nil(t, err)
    62  
    63  	// Test LoadOrStoreNewByFP
    64  	s, loaded, err = streams.LoadOrStoreNewByFP(ss[1].fp, func() (*stream, error) {
    65  		return ss[1], nil
    66  	}, nil)
    67  	require.Equal(t, s, ss[1])
    68  	require.False(t, loaded)
    69  	require.Nil(t, err)
    70  
    71  	require.Equal(t, len(ss), streams.Len())
    72  
    73  	for _, st := range ss {
    74  		s, loaded = streams.Load(st.labelsString)
    75  		require.Equal(t, st, s)
    76  		require.True(t, loaded)
    77  
    78  		s, loaded = streams.LoadByFP(st.fp)
    79  		require.Equal(t, st, s)
    80  		require.True(t, loaded)
    81  	}
    82  
    83  	// Test Delete
    84  	for _, st := range ss {
    85  		deleted := streams.Delete(st)
    86  		require.True(t, deleted)
    87  
    88  		s, loaded = streams.Load(st.labelsString)
    89  		require.Nil(t, s)
    90  		require.False(t, loaded)
    91  
    92  		s, loaded = streams.LoadByFP(st.fp)
    93  		require.Nil(t, s)
    94  		require.False(t, loaded)
    95  	}
    96  
    97  	require.Equal(t, 0, streams.Len())
    98  
    99  	// Test Store
   100  	streams.Store(ss[0].labelsString, ss[0])
   101  
   102  	s, loaded = streams.Load(ss[0].labelsString)
   103  	require.Equal(t, ss[0], s)
   104  	require.True(t, loaded)
   105  
   106  	s, loaded = streams.LoadByFP(ss[0].fp)
   107  	require.Equal(t, ss[0], s)
   108  	require.True(t, loaded)
   109  
   110  	// Test StoreByFP
   111  	streams.StoreByFP(ss[1].fp, ss[1])
   112  
   113  	s, loaded = streams.Load(ss[1].labelsString)
   114  	require.Equal(t, ss[1], s)
   115  	require.True(t, loaded)
   116  
   117  	s, loaded = streams.LoadByFP(ss[1].fp)
   118  	require.Equal(t, ss[1], s)
   119  	require.True(t, loaded)
   120  
   121  	require.Equal(t, len(ss), streams.Len())
   122  }