github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/ruler/base/manager_test.go (about)

     1  package base
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/go-kit/log"
     9  	"github.com/prometheus/client_golang/prometheus"
    10  	"github.com/prometheus/prometheus/model/labels"
    11  	"github.com/prometheus/prometheus/notifier"
    12  	promRules "github.com/prometheus/prometheus/rules"
    13  	"github.com/stretchr/testify/require"
    14  	"go.uber.org/atomic"
    15  
    16  	"github.com/grafana/loki/pkg/ruler/rulespb"
    17  	"github.com/grafana/loki/pkg/util/test"
    18  )
    19  
    20  func TestSyncRuleGroups(t *testing.T) {
    21  	dir := t.TempDir()
    22  
    23  	m, err := NewDefaultMultiTenantManager(Config{RulePath: dir}, factory, nil, log.NewNopLogger())
    24  	require.NoError(t, err)
    25  
    26  	const user = "testUser"
    27  
    28  	userRules := map[string]rulespb.RuleGroupList{
    29  		user: {
    30  			&rulespb.RuleGroupDesc{
    31  				Name:      "group1",
    32  				Namespace: "ns",
    33  				Interval:  1 * time.Minute,
    34  				User:      user,
    35  			},
    36  		},
    37  	}
    38  	m.SyncRuleGroups(context.Background(), userRules)
    39  
    40  	mgr := getManager(m, user)
    41  	require.NotNil(t, mgr)
    42  
    43  	test.Poll(t, 1*time.Second, true, func() interface{} {
    44  		return mgr.(*mockRulesManager).running.Load()
    45  	})
    46  
    47  	// Verify that user rule groups are now cached locally.
    48  	{
    49  		users, err := m.mapper.users()
    50  		require.NoError(t, err)
    51  		require.Equal(t, []string{user}, users)
    52  	}
    53  
    54  	// Passing empty map / nil stops all managers.
    55  	m.SyncRuleGroups(context.Background(), nil)
    56  	require.Nil(t, getManager(m, user))
    57  
    58  	// Make sure old manager was stopped.
    59  	test.Poll(t, 1*time.Second, false, func() interface{} {
    60  		return mgr.(*mockRulesManager).running.Load()
    61  	})
    62  
    63  	// Verify that local rule groups were removed.
    64  	{
    65  		users, err := m.mapper.users()
    66  		require.NoError(t, err)
    67  		require.Equal(t, []string(nil), users)
    68  	}
    69  
    70  	// Resync same rules as before. Previously this didn't restart the manager.
    71  	m.SyncRuleGroups(context.Background(), userRules)
    72  
    73  	newMgr := getManager(m, user)
    74  	require.NotNil(t, newMgr)
    75  	require.True(t, mgr != newMgr)
    76  
    77  	test.Poll(t, 1*time.Second, true, func() interface{} {
    78  		return newMgr.(*mockRulesManager).running.Load()
    79  	})
    80  
    81  	// Verify that user rule groups are cached locally again.
    82  	{
    83  		users, err := m.mapper.users()
    84  		require.NoError(t, err)
    85  		require.Equal(t, []string{user}, users)
    86  	}
    87  
    88  	m.Stop()
    89  
    90  	test.Poll(t, 1*time.Second, false, func() interface{} {
    91  		return newMgr.(*mockRulesManager).running.Load()
    92  	})
    93  }
    94  
    95  func getManager(m *DefaultMultiTenantManager, user string) RulesManager {
    96  	m.userManagerMtx.Lock()
    97  	defer m.userManagerMtx.Unlock()
    98  
    99  	return m.userManagers[user]
   100  }
   101  
   102  func factory(_ context.Context, _ string, _ *notifier.Manager, _ log.Logger, _ prometheus.Registerer) RulesManager {
   103  	return &mockRulesManager{done: make(chan struct{})}
   104  }
   105  
   106  type mockRulesManager struct {
   107  	running atomic.Bool
   108  	done    chan struct{}
   109  }
   110  
   111  func (m *mockRulesManager) Run() {
   112  	m.running.Store(true)
   113  	<-m.done
   114  }
   115  
   116  func (m *mockRulesManager) Stop() {
   117  	m.running.Store(false)
   118  	close(m.done)
   119  }
   120  
   121  func (m *mockRulesManager) Update(_ time.Duration, _ []string, _ labels.Labels, _ string, ruleGroupPostProcessFunc promRules.RuleGroupPostProcessFunc) error {
   122  	return nil
   123  }
   124  
   125  func (m *mockRulesManager) RuleGroups() []*promRules.Group {
   126  	return nil
   127  }