...

Source file src/github.com/prometheus/alertmanager/test/with_api_v1/mock.go

Documentation: github.com/prometheus/alertmanager/test/with_api_v1

     1  // Copyright 2015 Prometheus Team
     2  // Licensed under the Apache License, Version 2.0 (the "License");
     3  // you may not use this file except in compliance with the License.
     4  // You may obtain a copy of the License at
     5  //
     6  // http://www.apache.org/licenses/LICENSE-2.0
     7  //
     8  // Unless required by applicable law or agreed to in writing, software
     9  // distributed under the License is distributed on an "AS IS" BASIS,
    10  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package test
    15  
    16  import (
    17  	"encoding/json"
    18  	"fmt"
    19  	"net"
    20  	"net/http"
    21  	"reflect"
    22  	"sync"
    23  	"time"
    24  
    25  	"github.com/prometheus/common/model"
    26  
    27  	"github.com/prometheus/alertmanager/notify/webhook"
    28  	"github.com/prometheus/alertmanager/pkg/labels"
    29  	"github.com/prometheus/alertmanager/types"
    30  )
    31  
    32  // At is a convenience method to allow for declarative syntax of Acceptance
    33  // test definitions.
    34  func At(ts float64) float64 {
    35  	return ts
    36  }
    37  
    38  type Interval struct {
    39  	start, end float64
    40  }
    41  
    42  func (iv Interval) String() string {
    43  	return fmt.Sprintf("[%v,%v]", iv.start, iv.end)
    44  }
    45  
    46  func (iv Interval) contains(f float64) bool {
    47  	return f >= iv.start && f <= iv.end
    48  }
    49  
    50  // Between is a convenience constructor for an interval for declarative syntax
    51  // of Acceptance test definitions.
    52  func Between(start, end float64) Interval {
    53  	return Interval{start: start, end: end}
    54  }
    55  
    56  // TestSilence models a model.Silence with relative times.
    57  type TestSilence struct {
    58  	id               string
    59  	match            []string
    60  	matchRE          []string
    61  	startsAt, endsAt float64
    62  
    63  	mtx sync.RWMutex
    64  }
    65  
    66  // Silence creates a new TestSilence active for the relative interval given
    67  // by start and end.
    68  func Silence(start, end float64) *TestSilence {
    69  	return &TestSilence{
    70  		startsAt: start,
    71  		endsAt:   end,
    72  	}
    73  }
    74  
    75  // Match adds a new plain matcher to the silence.
    76  func (s *TestSilence) Match(v ...string) *TestSilence {
    77  	s.match = append(s.match, v...)
    78  	return s
    79  }
    80  
    81  // MatchRE adds a new regex matcher to the silence
    82  func (s *TestSilence) MatchRE(v ...string) *TestSilence {
    83  	if len(v)%2 == 1 {
    84  		panic("bad key/values")
    85  	}
    86  	s.matchRE = append(s.matchRE, v...)
    87  	return s
    88  }
    89  
    90  // SetID sets the silence ID.
    91  func (s *TestSilence) SetID(ID string) {
    92  	s.mtx.Lock()
    93  	defer s.mtx.Unlock()
    94  	s.id = ID
    95  }
    96  
    97  // ID gets the silence ID.
    98  func (s *TestSilence) ID() string {
    99  	s.mtx.RLock()
   100  	defer s.mtx.RUnlock()
   101  	return s.id
   102  }
   103  
   104  // nativeSilence converts the declared test silence into a regular
   105  // silence with resolved times.
   106  func (s *TestSilence) nativeSilence(opts *AcceptanceOpts) *types.Silence {
   107  	nsil := &types.Silence{}
   108  
   109  	for i := 0; i < len(s.match); i += 2 {
   110  		nsil.Matchers = append(nsil.Matchers, &labels.Matcher{
   111  			Type:  labels.MatchEqual,
   112  			Name:  s.match[i],
   113  			Value: s.match[i+1],
   114  		})
   115  	}
   116  	for i := 0; i < len(s.matchRE); i += 2 {
   117  		m, err := labels.NewMatcher(labels.MatchRegexp, s.matchRE[i], s.matchRE[i+1])
   118  		if err != nil {
   119  			panic(err)
   120  		}
   121  		nsil.Matchers = append(nsil.Matchers, m)
   122  	}
   123  
   124  	if s.startsAt > 0 {
   125  		nsil.StartsAt = opts.expandTime(s.startsAt)
   126  	}
   127  	if s.endsAt > 0 {
   128  		nsil.EndsAt = opts.expandTime(s.endsAt)
   129  	}
   130  	nsil.Comment = "some comment"
   131  	nsil.CreatedBy = "admin@example.com"
   132  
   133  	return nsil
   134  }
   135  
   136  // TestAlert models a model.Alert with relative times.
   137  type TestAlert struct {
   138  	labels           model.LabelSet
   139  	annotations      model.LabelSet
   140  	startsAt, endsAt float64
   141  }
   142  
   143  // Alert creates a new alert declaration with the given key/value pairs
   144  // as identifying labels.
   145  func Alert(keyval ...interface{}) *TestAlert {
   146  	if len(keyval)%2 == 1 {
   147  		panic("bad key/values")
   148  	}
   149  	a := &TestAlert{
   150  		labels:      model.LabelSet{},
   151  		annotations: model.LabelSet{},
   152  	}
   153  
   154  	for i := 0; i < len(keyval); i += 2 {
   155  		ln := model.LabelName(keyval[i].(string))
   156  		lv := model.LabelValue(keyval[i+1].(string))
   157  
   158  		a.labels[ln] = lv
   159  	}
   160  
   161  	return a
   162  }
   163  
   164  // nativeAlert converts the declared test alert into a full alert based
   165  // on the given parameters.
   166  func (a *TestAlert) nativeAlert(opts *AcceptanceOpts) *model.Alert {
   167  	na := &model.Alert{
   168  		Labels:      a.labels,
   169  		Annotations: a.annotations,
   170  	}
   171  
   172  	if a.startsAt > 0 {
   173  		na.StartsAt = opts.expandTime(a.startsAt)
   174  	}
   175  	if a.endsAt > 0 {
   176  		na.EndsAt = opts.expandTime(a.endsAt)
   177  	}
   178  	return na
   179  }
   180  
   181  // Annotate the alert with the given key/value pairs.
   182  func (a *TestAlert) Annotate(keyval ...interface{}) *TestAlert {
   183  	if len(keyval)%2 == 1 {
   184  		panic("bad key/values")
   185  	}
   186  
   187  	for i := 0; i < len(keyval); i += 2 {
   188  		ln := model.LabelName(keyval[i].(string))
   189  		lv := model.LabelValue(keyval[i+1].(string))
   190  
   191  		a.annotations[ln] = lv
   192  	}
   193  
   194  	return a
   195  }
   196  
   197  // Active declares the relative activity time for this alert. It
   198  // must be a single starting value or two values where the second value
   199  // declares the resolved time.
   200  func (a *TestAlert) Active(tss ...float64) *TestAlert {
   201  	if len(tss) > 2 || len(tss) == 0 {
   202  		panic("only one or two timestamps allowed")
   203  	}
   204  	if len(tss) == 2 {
   205  		a.endsAt = tss[1]
   206  	}
   207  	a.startsAt = tss[0]
   208  
   209  	return a
   210  }
   211  
   212  func equalAlerts(a, b *model.Alert, opts *AcceptanceOpts) bool {
   213  	if !reflect.DeepEqual(a.Labels, b.Labels) {
   214  		return false
   215  	}
   216  	if !reflect.DeepEqual(a.Annotations, b.Annotations) {
   217  		return false
   218  	}
   219  
   220  	if !equalTime(a.StartsAt, b.StartsAt, opts) {
   221  		return false
   222  	}
   223  	if !equalTime(a.EndsAt, b.EndsAt, opts) {
   224  		return false
   225  	}
   226  	return true
   227  }
   228  
   229  func equalTime(a, b time.Time, opts *AcceptanceOpts) bool {
   230  	if a.IsZero() != b.IsZero() {
   231  		return false
   232  	}
   233  
   234  	diff := a.Sub(b)
   235  	if diff < 0 {
   236  		diff = -diff
   237  	}
   238  	return diff <= opts.Tolerance
   239  }
   240  
   241  type MockWebhook struct {
   242  	opts      *AcceptanceOpts
   243  	collector *Collector
   244  	listener  net.Listener
   245  
   246  	Func func(timestamp float64) bool
   247  }
   248  
   249  func NewWebhook(c *Collector) *MockWebhook {
   250  	l, err := net.Listen("tcp4", "localhost:0")
   251  	if err != nil {
   252  		// TODO(fabxc): if shutdown of mock destinations ever becomes a concern
   253  		// we want to shut them down after test completion. Then we might want to
   254  		// log the error properly, too.
   255  		panic(err)
   256  	}
   257  	wh := &MockWebhook{
   258  		listener:  l,
   259  		collector: c,
   260  		opts:      c.opts,
   261  	}
   262  	go func() {
   263  		if err := http.Serve(l, wh); err != nil {
   264  			panic(err)
   265  		}
   266  	}()
   267  
   268  	return wh
   269  }
   270  
   271  func (ws *MockWebhook) ServeHTTP(w http.ResponseWriter, req *http.Request) {
   272  	// Inject Func if it exists.
   273  	if ws.Func != nil {
   274  		if ws.Func(ws.opts.relativeTime(time.Now())) {
   275  			return
   276  		}
   277  	}
   278  
   279  	dec := json.NewDecoder(req.Body)
   280  	defer req.Body.Close()
   281  
   282  	var v webhook.Message
   283  	if err := dec.Decode(&v); err != nil {
   284  		panic(err)
   285  	}
   286  
   287  	// Transform the webhook message alerts back into model.Alerts.
   288  	var alerts model.Alerts
   289  	for _, a := range v.Alerts {
   290  		var (
   291  			labels      = model.LabelSet{}
   292  			annotations = model.LabelSet{}
   293  		)
   294  		for k, v := range a.Labels {
   295  			labels[model.LabelName(k)] = model.LabelValue(v)
   296  		}
   297  		for k, v := range a.Annotations {
   298  			annotations[model.LabelName(k)] = model.LabelValue(v)
   299  		}
   300  
   301  		alerts = append(alerts, &model.Alert{
   302  			Labels:       labels,
   303  			Annotations:  annotations,
   304  			StartsAt:     a.StartsAt,
   305  			EndsAt:       a.EndsAt,
   306  			GeneratorURL: a.GeneratorURL,
   307  		})
   308  	}
   309  
   310  	ws.collector.add(alerts...)
   311  }
   312  
   313  func (ws *MockWebhook) Address() string {
   314  	return ws.listener.Addr().String()
   315  }
   316  

View as plain text