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

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