github.com/grafana/pyroscope@v1.18.0/pkg/metastore/client/client_test.go (about)

     1  package metastoreclient
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"testing"
     7  
     8  	"github.com/grafana/dskit/flagext"
     9  	"github.com/grafana/dskit/grpcclient"
    10  	"github.com/stretchr/testify/mock"
    11  	"github.com/stretchr/testify/require"
    12  
    13  	metastorev1 "github.com/grafana/pyroscope/api/gen/proto/go/metastore/v1"
    14  	"github.com/grafana/pyroscope/pkg/test"
    15  	"github.com/grafana/pyroscope/pkg/test/mocks/mockdiscovery"
    16  )
    17  
    18  func TestUnavailable(t *testing.T) {
    19  	d := mockdiscovery.NewMockDiscovery(t)
    20  	d.On("Subscribe", mock.Anything).Return()
    21  	l := test.NewTestingLogger(t)
    22  	c := New(l, grpcclient.Config{}, d)
    23  
    24  	d.On("Rediscover").Run(func(args mock.Arguments) {
    25  	}).Return()
    26  
    27  	c.updateServers(createServers([]int{30030, 30031, 30032}))
    28  	res, err := c.AddBlock(context.Background(), &metastorev1.AddBlockRequest{})
    29  	require.Error(t, err)
    30  	require.Nil(t, res)
    31  
    32  }
    33  
    34  func TestUnavailable_Rediscover_Wrong_Leader(t *testing.T) {
    35  	t.Run("AddBlock", func(t *testing.T) {
    36  		testRediscoverWrongLeader(t, func(c *Client) {
    37  			res, err := c.AddBlock(context.Background(), &metastorev1.AddBlockRequest{})
    38  			require.NoError(t, err)
    39  			require.NotNil(t, res)
    40  		})
    41  	})
    42  	t.Run("QueryMetadata", func(t *testing.T) {
    43  		testRediscoverWrongLeader(t, func(c *Client) {
    44  			res, err := c.QueryMetadata(context.Background(), &metastorev1.QueryMetadataRequest{})
    45  			require.NoError(t, err)
    46  			require.NotNil(t, res)
    47  		})
    48  	})
    49  	t.Run("PollCompactionJobs", func(t *testing.T) {
    50  		testRediscoverWrongLeader(t, func(c *Client) {
    51  			res, err := c.PollCompactionJobs(context.Background(), &metastorev1.PollCompactionJobsRequest{})
    52  			require.NoError(t, err)
    53  			require.NotNil(t, res)
    54  		})
    55  	})
    56  	t.Run("GetProfileStats", func(t *testing.T) {
    57  		testRediscoverWrongLeader(t, func(c *Client) {
    58  			res, err := c.GetTenant(context.Background(), &metastorev1.GetTenantRequest{})
    59  			require.NoError(t, err)
    60  			require.NotNil(t, res)
    61  		})
    62  	})
    63  }
    64  
    65  func testRediscoverWrongLeader(t *testing.T, f func(c *Client)) {
    66  	d := mockdiscovery.NewMockDiscovery(t)
    67  	d.On("Subscribe", mock.Anything).Return()
    68  	l := test.NewTestingLogger(t)
    69  	config := &grpcclient.Config{}
    70  	flagext.DefaultValues(config)
    71  
    72  	dServers1 := createServers([]int{30031, 30032, 30033})
    73  
    74  	dServers2 := createServers([]int{40031, 40032, 40033})
    75  	mockServers2, dialOpts := createMockServers(t, l, dServers2)
    76  	defer mockServers2.Close()
    77  
    78  	c := New(l, *config, d, dialOpts...)
    79  	m := sync.Mutex{}
    80  	verify := func() {}
    81  	initWrongLeaderCalled := false
    82  	d.On("Rediscover", mock.Anything).Run(func(args mock.Arguments) {
    83  		m.Lock()
    84  		defer m.Unlock()
    85  		if !initWrongLeaderCalled {
    86  			initWrongLeaderCalled = true
    87  			verify = mockServers2.InitWrongLeader()
    88  			// call updateServers twice
    89  			c.updateServers(dServers2)
    90  			c.updateServers(dServers2)
    91  		}
    92  	}).Return()
    93  
    94  	c.updateServers(dServers1)
    95  	f(c)
    96  	verify()
    97  }
    98  
    99  func TestServerError(t *testing.T) {
   100  	d := mockdiscovery.NewMockDiscovery(t)
   101  	d.On("Subscribe", mock.Anything).Return()
   102  	l := test.NewTestingLogger(t)
   103  	c := New(l, grpcclient.Config{}, d)
   104  
   105  	d.On("Rediscover").Run(func(args mock.Arguments) {
   106  	}).Return()
   107  
   108  	res, err := c.AddBlock(context.Background(), &metastorev1.AddBlockRequest{})
   109  	require.Error(t, err)
   110  	require.Nil(t, res)
   111  }