...

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

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

     1  // Copyright 2018 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/http"
    20  	"net/http/httptest"
    21  	"reflect"
    22  	"sync"
    23  	"testing"
    24  	"time"
    25  
    26  	"github.com/go-openapi/strfmt"
    27  
    28  	"github.com/prometheus/alertmanager/api/v2/models"
    29  	"github.com/prometheus/alertmanager/notify/webhook"
    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) *models.Silence {
   107  	nsil := &models.Silence{}
   108  
   109  	t := false
   110  	for i := 0; i < len(s.match); i += 2 {
   111  		nsil.Matchers = append(nsil.Matchers, &models.Matcher{
   112  			Name:    &s.match[i],
   113  			Value:   &s.match[i+1],
   114  			IsRegex: &t,
   115  		})
   116  	}
   117  	t = true
   118  	for i := 0; i < len(s.matchRE); i += 2 {
   119  		nsil.Matchers = append(nsil.Matchers, &models.Matcher{
   120  			Name:    &s.matchRE[i],
   121  			Value:   &s.matchRE[i+1],
   122  			IsRegex: &t,
   123  		})
   124  	}
   125  
   126  	if s.startsAt > 0 {
   127  		start := strfmt.DateTime(opts.expandTime(s.startsAt))
   128  		nsil.StartsAt = &start
   129  	}
   130  	if s.endsAt > 0 {
   131  		end := strfmt.DateTime(opts.expandTime(s.endsAt))
   132  		nsil.EndsAt = &end
   133  	}
   134  	comment := "some comment"
   135  	createdBy := "admin@example.com"
   136  	nsil.Comment = &comment
   137  	nsil.CreatedBy = &createdBy
   138  
   139  	return nsil
   140  }
   141  
   142  // TestAlert models a model.Alert with relative times.
   143  type TestAlert struct {
   144  	labels           models.LabelSet
   145  	annotations      models.LabelSet
   146  	startsAt, endsAt float64
   147  }
   148  
   149  // Alert creates a new alert declaration with the given key/value pairs
   150  // as identifying labels.
   151  func Alert(keyval ...interface{}) *TestAlert {
   152  	if len(keyval)%2 == 1 {
   153  		panic("bad key/values")
   154  	}
   155  	a := &TestAlert{
   156  		labels:      models.LabelSet{},
   157  		annotations: models.LabelSet{},
   158  	}
   159  
   160  	for i := 0; i < len(keyval); i += 2 {
   161  		ln := keyval[i].(string)
   162  		lv := keyval[i+1].(string)
   163  
   164  		a.labels[ln] = lv
   165  	}
   166  
   167  	return a
   168  }
   169  
   170  // nativeAlert converts the declared test alert into a full alert based
   171  // on the given parameters.
   172  func (a *TestAlert) nativeAlert(opts *AcceptanceOpts) *models.GettableAlert {
   173  	na := &models.GettableAlert{
   174  		Alert: models.Alert{
   175  			Labels: a.labels,
   176  		},
   177  		Annotations: a.annotations,
   178  		StartsAt:    &strfmt.DateTime{},
   179  		EndsAt:      &strfmt.DateTime{},
   180  	}
   181  
   182  	if a.startsAt > 0 {
   183  		start := strfmt.DateTime(opts.expandTime(a.startsAt))
   184  		na.StartsAt = &start
   185  	}
   186  	if a.endsAt > 0 {
   187  		end := strfmt.DateTime(opts.expandTime(a.endsAt))
   188  		na.EndsAt = &end
   189  	}
   190  
   191  	return na
   192  }
   193  
   194  // Annotate the alert with the given key/value pairs.
   195  func (a *TestAlert) Annotate(keyval ...interface{}) *TestAlert {
   196  	if len(keyval)%2 == 1 {
   197  		panic("bad key/values")
   198  	}
   199  
   200  	for i := 0; i < len(keyval); i += 2 {
   201  		ln := keyval[i].(string)
   202  		lv := keyval[i+1].(string)
   203  
   204  		a.annotations[ln] = lv
   205  	}
   206  
   207  	return a
   208  }
   209  
   210  // Active declares the relative activity time for this alert. It
   211  // must be a single starting value or two values where the second value
   212  // declares the resolved time.
   213  func (a *TestAlert) Active(tss ...float64) *TestAlert {
   214  	if len(tss) > 2 || len(tss) == 0 {
   215  		panic("only one or two timestamps allowed")
   216  	}
   217  	if len(tss) == 2 {
   218  		a.endsAt = tss[1]
   219  	}
   220  	a.startsAt = tss[0]
   221  
   222  	return a
   223  }
   224  
   225  func equalAlerts(a, b *models.GettableAlert, opts *AcceptanceOpts) bool {
   226  	if !reflect.DeepEqual(a.Labels, b.Labels) {
   227  		return false
   228  	}
   229  	if !reflect.DeepEqual(a.Annotations, b.Annotations) {
   230  		return false
   231  	}
   232  
   233  	if !equalTime(time.Time(*a.StartsAt), time.Time(*b.StartsAt), opts) {
   234  		return false
   235  	}
   236  	if (a.EndsAt == nil) != (b.EndsAt == nil) {
   237  		return false
   238  	}
   239  	if !(a.EndsAt == nil) && !(b.EndsAt == nil) && !equalTime(time.Time(*a.EndsAt), time.Time(*b.EndsAt), opts) {
   240  		return false
   241  	}
   242  	return true
   243  }
   244  
   245  func equalTime(a, b time.Time, opts *AcceptanceOpts) bool {
   246  	if a.IsZero() != b.IsZero() {
   247  		return false
   248  	}
   249  
   250  	diff := a.Sub(b)
   251  	if diff < 0 {
   252  		diff = -diff
   253  	}
   254  	return diff <= opts.Tolerance
   255  }
   256  
   257  type MockWebhook struct {
   258  	opts      *AcceptanceOpts
   259  	collector *Collector
   260  	addr      string
   261  
   262  	// Func is called early on when retrieving a notification by an
   263  	// Alertmanager. If Func returns true, the given notification is dropped.
   264  	// See sample usage in `send_test.go/TestRetry()`.
   265  	Func func(timestamp float64) bool
   266  }
   267  
   268  func NewWebhook(t *testing.T, c *Collector) *MockWebhook {
   269  	t.Helper()
   270  
   271  	wh := &MockWebhook{
   272  		collector: c,
   273  		opts:      c.opts,
   274  	}
   275  
   276  	server := httptest.NewServer(wh)
   277  	wh.addr = server.Listener.Addr().String()
   278  
   279  	t.Cleanup(func() {
   280  		server.Close()
   281  	})
   282  
   283  	return wh
   284  }
   285  
   286  func (ws *MockWebhook) ServeHTTP(w http.ResponseWriter, req *http.Request) {
   287  	// Inject Func if it exists.
   288  	if ws.Func != nil {
   289  		if ws.Func(ws.opts.relativeTime(time.Now())) {
   290  			return
   291  		}
   292  	}
   293  
   294  	dec := json.NewDecoder(req.Body)
   295  	defer req.Body.Close()
   296  
   297  	var v webhook.Message
   298  	if err := dec.Decode(&v); err != nil {
   299  		panic(err)
   300  	}
   301  
   302  	// Transform the webhook message alerts back into model.Alerts.
   303  	var alerts models.GettableAlerts
   304  	for _, a := range v.Alerts {
   305  		var (
   306  			labels      = models.LabelSet{}
   307  			annotations = models.LabelSet{}
   308  		)
   309  		for k, v := range a.Labels {
   310  			labels[k] = v
   311  		}
   312  		for k, v := range a.Annotations {
   313  			annotations[k] = v
   314  		}
   315  
   316  		start := strfmt.DateTime(a.StartsAt)
   317  		end := strfmt.DateTime(a.EndsAt)
   318  
   319  		alerts = append(alerts, &models.GettableAlert{
   320  			Alert: models.Alert{
   321  				Labels:       labels,
   322  				GeneratorURL: strfmt.URI(a.GeneratorURL),
   323  			},
   324  			Annotations: annotations,
   325  			StartsAt:    &start,
   326  			EndsAt:      &end,
   327  		})
   328  	}
   329  
   330  	ws.collector.add(alerts...)
   331  }
   332  
   333  func (ws *MockWebhook) Address() string {
   334  	return ws.addr
   335  }
   336  

View as plain text