github.com/sequix/cortex@v1.1.6/pkg/ruler/ruler_test.go (about)

     1  package ruler
     2  
     3  import (
     4  	"context"
     5  	"net/http"
     6  	"net/http/httptest"
     7  	"sync"
     8  	"testing"
     9  	"time"
    10  
    11  	"github.com/prometheus/prometheus/pkg/labels"
    12  	"github.com/prometheus/prometheus/promql"
    13  
    14  	"github.com/sequix/cortex/pkg/configs"
    15  	client_config "github.com/sequix/cortex/pkg/configs/client"
    16  	"github.com/sequix/cortex/pkg/querier"
    17  	"github.com/sequix/cortex/pkg/ring"
    18  	"github.com/sequix/cortex/pkg/ring/kv/codec"
    19  	"github.com/sequix/cortex/pkg/ring/kv/consul"
    20  	"github.com/sequix/cortex/pkg/util/flagext"
    21  	"github.com/prometheus/prometheus/notifier"
    22  	"github.com/stretchr/testify/assert"
    23  	"github.com/weaveworks/common/user"
    24  )
    25  
    26  type mockRuleStore struct{}
    27  
    28  func (m *mockRuleStore) GetRules(ctx context.Context, since configs.ID) (map[string]configs.VersionedRulesConfig, error) {
    29  	return map[string]configs.VersionedRulesConfig{}, nil
    30  }
    31  
    32  func (m *mockRuleStore) GetAlerts(ctx context.Context, since configs.ID) (*client_config.ConfigsResponse, error) {
    33  	return nil, nil
    34  }
    35  
    36  func defaultRulerConfig() Config {
    37  	codec := codec.Proto{Factory: ring.ProtoDescFactory}
    38  	consul := consul.NewInMemoryClient(codec)
    39  	cfg := Config{}
    40  	flagext.DefaultValues(&cfg)
    41  	flagext.DefaultValues(&cfg.LifecyclerConfig)
    42  	cfg.LifecyclerConfig.RingConfig.ReplicationFactor = 1
    43  	cfg.LifecyclerConfig.RingConfig.KVStore.Mock = consul
    44  	cfg.LifecyclerConfig.NumTokens = 1
    45  	cfg.LifecyclerConfig.FinalSleep = time.Duration(0)
    46  	cfg.LifecyclerConfig.ListenPort = func(i int) *int { return &i }(0)
    47  	cfg.LifecyclerConfig.Addr = "localhost"
    48  	cfg.LifecyclerConfig.ID = "localhost"
    49  	return cfg
    50  }
    51  
    52  func newTestRuler(t *testing.T, cfg Config) *Ruler {
    53  	// TODO: Populate distributor and chunk store arguments to enable
    54  	// other kinds of tests.
    55  
    56  	engine := promql.NewEngine(promql.EngineOpts{
    57  		MaxSamples:    1e6,
    58  		MaxConcurrent: 20,
    59  		Timeout:       2 * time.Minute,
    60  	})
    61  	queryable := querier.NewQueryable(nil, nil, nil, 0)
    62  	ruler, err := NewRuler(cfg, engine, queryable, nil, &mockRuleStore{})
    63  	if err != nil {
    64  		t.Fatal(err)
    65  	}
    66  	return ruler
    67  }
    68  
    69  func TestNotifierSendsUserIDHeader(t *testing.T) {
    70  	var wg sync.WaitGroup
    71  	wg.Add(1) // We want one request to our test HTTP server.
    72  	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    73  		userID, _, err := user.ExtractOrgIDFromHTTPRequest(r)
    74  		assert.NoError(t, err)
    75  		assert.Equal(t, userID, "1")
    76  		wg.Done()
    77  	}))
    78  	defer ts.Close()
    79  
    80  	cfg := defaultRulerConfig()
    81  	cfg.AlertmanagerURL.Set(ts.URL)
    82  	cfg.AlertmanagerDiscovery = false
    83  
    84  	r := newTestRuler(t, cfg)
    85  	n, err := r.getOrCreateNotifier("1")
    86  	if err != nil {
    87  		t.Fatal(err)
    88  	}
    89  	for _, not := range r.notifiers {
    90  		defer not.stop()
    91  	}
    92  	// Loop until notifier discovery syncs up
    93  	for len(n.Alertmanagers()) == 0 {
    94  		time.Sleep(10 * time.Millisecond)
    95  	}
    96  	n.Send(&notifier.Alert{
    97  		Labels: labels.Labels{labels.Label{Name: "alertname", Value: "testalert"}},
    98  	})
    99  
   100  	wg.Wait()
   101  }