diff --git a/rules/manager_test.go b/rules/manager_test.go index db2e2b468..5d6de382f 100644 --- a/rules/manager_test.go +++ b/rules/manager_test.go @@ -817,14 +817,14 @@ func TestUpdateSetsSourceTenants(t *testing.T) { defer ruleManager.Stop() rgs, errs := rulefmt.ParseFile("fixtures/rules_with_source_tenants.yaml") - require.Equal(t, 0, len(errs), "file parsing failures") + require.Empty(t, errs, "file parsing failures") tmpFile, err := ioutil.TempFile("", "rules.test.*.yaml") require.NoError(t, err) defer os.Remove(tmpFile.Name()) defer tmpFile.Close() - reloadRules(rgs, t, tmpFile, ruleManager) + reloadRules(rgs, t, tmpFile, ruleManager, 0) // check that all source tenants were actually set require.Len(t, ruleManager.groups, len(rgs.Groups)) @@ -837,6 +837,55 @@ func TestUpdateSetsSourceTenants(t *testing.T) { } } +func TestGroupEvaluationContextFuncIsCalledWhenSupplied(t *testing.T) { + type testContextKeyType string + var testContextKey testContextKeyType = "TestGroupEvaluationContextFuncIsCalledWhenSupplied" + oldContextTestValue := context.Background().Value(testContextKey) + + contextTestValueChannel := make(chan interface{}) + mockQueryFunc := func(ctx context.Context, qs string, t time.Time) (promql.Vector, error) { + contextTestValueChannel <- ctx.Value(testContextKey) + return promql.Vector{}, nil + } + + mockContextWrapFunc := func(ctx context.Context, g *Group) context.Context { + return context.WithValue(ctx, testContextKey, 42) + } + + st := teststorage.New(t) + defer st.Close() + + ruleManager := NewManager(&ManagerOptions{ + Appendable: st, + Queryable: st, + QueryFunc: mockQueryFunc, + Context: context.Background(), + Logger: log.NewNopLogger(), + GroupEvaluationContextFunc: mockContextWrapFunc, + }) + + rgs, errs := rulefmt.ParseFile("fixtures/rules_with_source_tenants.yaml") + require.Empty(t, errs, "file parsing failures") + + tmpFile, err := ioutil.TempFile("", "rules.test.*.yaml") + require.NoError(t, err) + defer os.Remove(tmpFile.Name()) + defer tmpFile.Close() + + // no filesystem is harmed when running this test, set the interval low + reloadRules(rgs, t, tmpFile, ruleManager, 10*time.Millisecond) + + ruleManager.start() + defer ruleManager.Stop() + + // check that all source tenants were actually set + require.Len(t, ruleManager.groups, len(rgs.Groups)) + + require.Nil(t, oldContextTestValue, "Context contained test key before the test, impossible") + newContextTestValue := <-contextTestValueChannel + require.Equal(t, 42, newContextTestValue, "Context does not contain the correct value that should be injected") +} + // ruleGroupsTest for running tests over rules. type ruleGroupsTest struct { Groups []ruleGroupTest `yaml:"groups"` @@ -879,18 +928,22 @@ func formatRules(r *rulefmt.RuleGroups) ruleGroupsTest { } } -func reloadRules(rgs *rulefmt.RuleGroups, t *testing.T, tmpFile *os.File, ruleManager *Manager) { +func reloadRules(rgs *rulefmt.RuleGroups, t *testing.T, tmpFile *os.File, ruleManager *Manager, interval time.Duration) { + if interval == 0 { + interval = 10 * time.Second + } + bs, err := yaml.Marshal(formatRules(rgs)) require.NoError(t, err) _, _ = tmpFile.Seek(0, 0) _, err = tmpFile.Write(bs) require.NoError(t, err) - err = ruleManager.Update(10*time.Second, []string{tmpFile.Name()}, nil, "") + err = ruleManager.Update(interval, []string{tmpFile.Name()}, nil, "") require.NoError(t, err) } func reloadAndValidate(rgs *rulefmt.RuleGroups, t *testing.T, tmpFile *os.File, ruleManager *Manager, expected map[string]labels.Labels, ogs map[string]*Group) { - reloadRules(rgs, t, tmpFile, ruleManager) + reloadRules(rgs, t, tmpFile, ruleManager, 0) for h, g := range ruleManager.groups { if ogs[h] == g {