github.com/muhammadn/cortex@v1.9.1-0.20220510110439-46bb7000d03d/pkg/alertmanager/state_persister_test.go (about)

     1  package alertmanager
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/go-kit/log"
    10  	"github.com/grafana/dskit/services"
    11  	"github.com/prometheus/alertmanager/cluster/clusterpb"
    12  	"github.com/stretchr/testify/assert"
    13  	"github.com/stretchr/testify/require"
    14  
    15  	"github.com/cortexproject/cortex/pkg/alertmanager/alertspb"
    16  	"github.com/cortexproject/cortex/pkg/alertmanager/alertstore"
    17  )
    18  
    19  type fakePersistableState struct {
    20  	PersistableState
    21  
    22  	position int
    23  	readyc   chan struct{}
    24  
    25  	getResult *clusterpb.FullState
    26  	getError  error
    27  }
    28  
    29  func (f *fakePersistableState) Position() int {
    30  	return f.position
    31  }
    32  
    33  func (f *fakePersistableState) GetFullState() (*clusterpb.FullState, error) {
    34  	return f.getResult, f.getError
    35  }
    36  
    37  func newFakePersistableState() *fakePersistableState {
    38  	return &fakePersistableState{
    39  		readyc: make(chan struct{}),
    40  	}
    41  }
    42  
    43  func (f *fakePersistableState) WaitReady(ctx context.Context) error {
    44  	<-f.readyc
    45  	return nil
    46  }
    47  
    48  type fakeStoreWrite struct {
    49  	user string
    50  	desc alertspb.FullStateDesc
    51  }
    52  
    53  type fakeStore struct {
    54  	alertstore.AlertStore
    55  
    56  	writesMtx sync.Mutex
    57  	writes    []fakeStoreWrite
    58  }
    59  
    60  func (f *fakeStore) SetFullState(ctx context.Context, user string, desc alertspb.FullStateDesc) error {
    61  	f.writesMtx.Lock()
    62  	defer f.writesMtx.Unlock()
    63  	f.writes = append(f.writes, fakeStoreWrite{user, desc})
    64  	return nil
    65  }
    66  
    67  func (f *fakeStore) getWrites() []fakeStoreWrite {
    68  	f.writesMtx.Lock()
    69  	defer f.writesMtx.Unlock()
    70  	return f.writes
    71  }
    72  
    73  func makeTestFullState() *clusterpb.FullState {
    74  	return &clusterpb.FullState{
    75  		Parts: []clusterpb.Part{
    76  			{
    77  				Key:  "key",
    78  				Data: []byte("data"),
    79  			},
    80  		},
    81  	}
    82  }
    83  
    84  func makeTestStatePersister(t *testing.T, position int, userID string) (*fakePersistableState, *fakeStore, *statePersister) {
    85  	state := newFakePersistableState()
    86  	state.position = position
    87  	store := &fakeStore{}
    88  	cfg := PersisterConfig{Interval: 1 * time.Second}
    89  
    90  	s := newStatePersister(cfg, userID, state, store, log.NewNopLogger(), nil)
    91  
    92  	require.NoError(t, s.StartAsync(context.Background()))
    93  	t.Cleanup(func() {
    94  		require.NoError(t, services.StopAndAwaitTerminated(context.Background(), s))
    95  	})
    96  
    97  	return state, store, s
    98  }
    99  
   100  func TestStatePersister_Position0ShouldWrite(t *testing.T) {
   101  	userID := "user-1"
   102  	state, store, s := makeTestStatePersister(t, 0, userID)
   103  
   104  	// Should not start until the state becomes ready.
   105  	{
   106  		time.Sleep(5 * time.Second)
   107  
   108  		assert.Equal(t, services.Starting, s.Service.State())
   109  		assert.Equal(t, 0, len(store.getWrites()))
   110  	}
   111  
   112  	// Should start successfully once the state returns from WaitReady.
   113  	{
   114  		state.getResult = makeTestFullState()
   115  		close(state.readyc)
   116  
   117  		assert.NoError(t, s.AwaitRunning(context.Background()))
   118  	}
   119  
   120  	// Should receive a write to the store.
   121  	{
   122  		var storeWrites []fakeStoreWrite
   123  		require.Eventually(t, func() bool {
   124  			storeWrites = store.getWrites()
   125  			return len(storeWrites) == 1
   126  		}, 5*time.Second, 100*time.Millisecond)
   127  
   128  		expectedDesc := alertspb.FullStateDesc{
   129  			State: makeTestFullState(),
   130  		}
   131  
   132  		assert.Equal(t, userID, storeWrites[0].user)
   133  		assert.Equal(t, expectedDesc, storeWrites[0].desc)
   134  	}
   135  }
   136  
   137  func TestStatePersister_Position1ShouldNotWrite(t *testing.T) {
   138  	state, store, s := makeTestStatePersister(t, 1, "x")
   139  
   140  	// Start the persister.
   141  	{
   142  		require.Equal(t, services.Starting, s.Service.State())
   143  
   144  		state.getResult = makeTestFullState()
   145  		close(state.readyc)
   146  
   147  		require.NoError(t, s.AwaitRunning(context.Background()))
   148  		require.Equal(t, services.Running, s.Service.State())
   149  	}
   150  
   151  	// Should not have stored anything, having passed the interval multiple times.
   152  	{
   153  		time.Sleep(5 * time.Second)
   154  
   155  		assert.Equal(t, 0, len(store.getWrites()))
   156  	}
   157  }