telegraf/plugins/outputs/cloud_pubsub/topic_stubbed.go

213 lines
4.5 KiB
Go
Raw Normal View History

package cloud_pubsub
import (
"context"
"errors"
"fmt"
"runtime"
"sync"
"testing"
"time"
"cloud.google.com/go/pubsub"
"encoding/base64"
"github.com/influxdata/telegraf"
"github.com/influxdata/telegraf/internal"
"github.com/influxdata/telegraf/plugins/parsers"
"github.com/influxdata/telegraf/plugins/serializers"
"google.golang.org/api/support/bundler"
)
const (
errMockFail = "this is an error"
)
type (
testMetric struct {
m telegraf.Metric
returnErr bool
}
bundledMsg struct {
*pubsub.Message
*stubResult
}
stubResult struct {
metricIds []string
sendError bool
err chan error
done chan struct{}
}
stubTopic struct {
Settings pubsub.PublishSettings
ReturnErr map[string]bool
parsers.Parser
*testing.T
stopped bool
pLock sync.Mutex
published map[string]*pubsub.Message
bundler *bundler.Bundler
bLock sync.Mutex
bundleCount int
}
)
func getTestResources(tT *testing.T, settings pubsub.PublishSettings, testM []testMetric) (*PubSub, *stubTopic, []telegraf.Metric) {
s, _ := serializers.NewInfluxSerializer()
metrics := make([]telegraf.Metric, len(testM))
t := &stubTopic{
T: tT,
ReturnErr: make(map[string]bool),
published: make(map[string]*pubsub.Message),
}
for i, tm := range testM {
metrics[i] = tm.m
if tm.returnErr {
v, _ := tm.m.GetField("value")
t.ReturnErr[v.(string)] = true
}
}
ps := &PubSub{
Project: "test-project",
Topic: "test-topic",
stubTopic: func(string) topic { return t },
PublishCountThreshold: settings.CountThreshold,
PublishByteThreshold: settings.ByteThreshold,
PublishNumGoroutines: settings.NumGoroutines,
PublishTimeout: internal.Duration{Duration: settings.Timeout},
}
ps.SetSerializer(s)
return ps, t, metrics
}
func (t *stubTopic) ID() string {
return "test-topic"
}
func (t *stubTopic) Stop() {
t.pLock.Lock()
defer t.pLock.Unlock()
t.stopped = true
t.bundler.Flush()
}
func (t *stubTopic) Publish(ctx context.Context, msg *pubsub.Message) publishResult {
t.pLock.Lock()
defer t.pLock.Unlock()
if t.stopped || ctx.Err() != nil {
t.Fatalf("publish called after stop")
}
ids := t.parseIDs(msg)
r := &stubResult{
metricIds: ids,
err: make(chan error, 1),
done: make(chan struct{}, 1),
}
for _, id := range ids {
_, ok := t.ReturnErr[id]
r.sendError = r.sendError || ok
}
bundled := &bundledMsg{msg, r}
err := t.bundler.Add(bundled, len(msg.Data))
if err != nil {
t.Fatalf("unexpected error while adding to bundle: %v", err)
}
return r
}
func (t *stubTopic) PublishSettings() pubsub.PublishSettings {
return t.Settings
}
func (t *stubTopic) SetPublishSettings(settings pubsub.PublishSettings) {
t.Settings = settings
t.initBundler()
}
func (t *stubTopic) initBundler() *stubTopic {
t.bundler = bundler.NewBundler(&bundledMsg{}, t.sendBundle())
t.bundler.DelayThreshold = 10 * time.Second
t.bundler.BundleCountThreshold = t.Settings.CountThreshold
if t.bundler.BundleCountThreshold > pubsub.MaxPublishRequestCount {
t.bundler.BundleCountThreshold = pubsub.MaxPublishRequestCount
}
t.bundler.BundleByteThreshold = t.Settings.ByteThreshold
t.bundler.BundleByteLimit = pubsub.MaxPublishRequestBytes
t.bundler.HandlerLimit = 25 * runtime.GOMAXPROCS(0)
return t
}
func (t *stubTopic) sendBundle() func(items interface{}) {
return func(items interface{}) {
t.bLock.Lock()
defer t.bLock.Unlock()
bundled := items.([]*bundledMsg)
for _, msg := range bundled {
r := msg.stubResult
for _, id := range r.metricIds {
t.published[id] = msg.Message
}
if r.sendError {
r.err <- errors.New(errMockFail)
} else {
r.done <- struct{}{}
}
}
t.bundleCount++
}
}
func (t *stubTopic) parseIDs(msg *pubsub.Message) []string {
p, _ := parsers.NewInfluxParser()
metrics, err := p.Parse(msg.Data)
if err != nil {
// Just attempt to base64-decode first before returning error.
d, err := base64.StdEncoding.DecodeString(string(msg.Data))
if err != nil {
t.Errorf("unable to base64-decode potential test message: %v", err)
}
metrics, err = p.Parse(d)
if err != nil {
t.Fatalf("unexpected parsing error: %v", err)
}
}
ids := make([]string, len(metrics))
for i, met := range metrics {
id, _ := met.GetField("value")
ids[i] = id.(string)
}
return ids
}
func (r *stubResult) Get(ctx context.Context) (string, error) {
select {
case <-ctx.Done():
return "", ctx.Err()
case err := <-r.err:
return "", err
case <-r.done:
return fmt.Sprintf("id-%s", r.metricIds[0]), nil
}
}