prometheus/discovery/discovery_test.go

133 lines
3.1 KiB
Go

// Copyright 2016 The Prometheus Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package discovery
import (
"context"
"sync"
"sync/atomic"
"testing"
"github.com/prometheus/prometheus/config"
yaml "gopkg.in/yaml.v2"
)
func TestTargetSetRecreatesTargetGroupsEveryRun(t *testing.T) {
verifyPresence := func(tgroups map[string]*config.TargetGroup, name string, present bool) {
if _, ok := tgroups[name]; ok != present {
msg := ""
if !present {
msg = "not "
}
t.Fatalf("'%s' should %sbe present in TargetSet.tgroups: %s", name, msg, tgroups)
}
}
cfg := &config.ServiceDiscoveryConfig{}
sOne := `
static_configs:
- targets: ["foo:9090"]
- targets: ["bar:9090"]
`
if err := yaml.Unmarshal([]byte(sOne), cfg); err != nil {
t.Fatalf("Unable to load YAML config sOne: %s", err)
}
called := make(chan struct{})
ts := NewTargetSet(&mockSyncer{
sync: func([]*config.TargetGroup) { called <- struct{}{} },
})
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
go ts.Run(ctx)
ts.UpdateProviders(ProvidersFromConfig(*cfg, nil))
<-called
verifyPresence(ts.tgroups, "static/0/0", true)
verifyPresence(ts.tgroups, "static/0/1", true)
sTwo := `
static_configs:
- targets: ["foo:9090"]
`
if err := yaml.Unmarshal([]byte(sTwo), cfg); err != nil {
t.Fatalf("Unable to load YAML config sTwo: %s", err)
}
ts.UpdateProviders(ProvidersFromConfig(*cfg, nil))
<-called
verifyPresence(ts.tgroups, "static/0/0", true)
verifyPresence(ts.tgroups, "static/0/1", false)
}
type mockSyncer struct {
sync func(tgs []*config.TargetGroup)
}
func (s *mockSyncer) Sync(tgs []*config.TargetGroup) {
if s.sync != nil {
s.sync(tgs)
}
}
type mockTargetProvider struct {
callCount *uint32
}
func (tp mockTargetProvider) Run(ctx context.Context, up chan<- []*config.TargetGroup) {
atomic.AddUint32(tp.callCount, 1)
up <- []*config.TargetGroup{{Source: "dummySource"}}
}
func TestTargetSetRunsSameTargetProviderMultipleTimes(t *testing.T) {
var wg sync.WaitGroup
wg.Add(2)
ts1 := NewTargetSet(&mockSyncer{
sync: func([]*config.TargetGroup) { wg.Done() },
})
ts2 := NewTargetSet(&mockSyncer{
sync: func([]*config.TargetGroup) { wg.Done() },
})
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
tp := mockTargetProvider{}
var callCount uint32
tp.callCount = &callCount
targetProviders := map[string]TargetProvider{}
targetProviders["testProvider"] = tp
go ts1.Run(ctx)
go ts2.Run(ctx)
ts1.UpdateProviders(targetProviders)
ts2.UpdateProviders(targetProviders)
wg.Wait()
if callCount != 2 {
t.Errorf("Was expecting 2 calls received %v", callCount)
}
}