github.com/lyft/flytestdlib@v0.3.12-0.20210213045714-8cdd111ecda1/promutils/labeled/gauge_test.go (about)

     1  package labeled
     2  
     3  import (
     4  	"context"
     5  	"strings"
     6  	"testing"
     7  
     8  	"github.com/lyft/flytestdlib/contextutils"
     9  	"github.com/lyft/flytestdlib/promutils"
    10  	"github.com/prometheus/client_golang/prometheus/testutil"
    11  	"github.com/stretchr/testify/assert"
    12  )
    13  
    14  func TestLabeledGauge(t *testing.T) {
    15  	UnsetMetricKeys()
    16  	assert.NotPanics(t, func() {
    17  		SetMetricKeys(contextutils.ProjectKey, contextutils.DomainKey, contextutils.WorkflowIDKey, contextutils.TaskIDKey, contextutils.LaunchPlanIDKey)
    18  	})
    19  
    20  	scope := promutils.NewScope("testscope")
    21  	ctx := context.Background()
    22  	ctx = contextutils.WithProjectDomain(ctx, "flyte", "dev")
    23  	g := NewGauge("unittest", "some desc", scope)
    24  	assert.NotNil(t, g)
    25  
    26  	g.Inc(ctx)
    27  
    28  	const header = `
    29  		# HELP testscope:unittest some desc
    30          # TYPE testscope:unittest gauge
    31  	`
    32  	var expected = `
    33          testscope:unittest{domain="dev",lp="",project="flyte",task="",wf=""} 1
    34  	`
    35  	err := testutil.CollectAndCompare(g.GaugeVec, strings.NewReader(header+expected))
    36  	assert.NoError(t, err)
    37  
    38  	g.Set(ctx, 42)
    39  	expected = `
    40          testscope:unittest{domain="dev",lp="",project="flyte",task="",wf=""} 42
    41  	`
    42  	err = testutil.CollectAndCompare(g.GaugeVec, strings.NewReader(header+expected))
    43  	assert.NoError(t, err)
    44  
    45  	g.Add(ctx, 1)
    46  	expected = `
    47          testscope:unittest{domain="dev",lp="",project="flyte",task="",wf=""} 43
    48  	`
    49  	err = testutil.CollectAndCompare(g.GaugeVec, strings.NewReader(header+expected))
    50  	assert.NoError(t, err)
    51  
    52  	g.Dec(ctx)
    53  	expected = `
    54          testscope:unittest{domain="dev",lp="",project="flyte",task="",wf=""} 42
    55  	`
    56  	err = testutil.CollectAndCompare(g.GaugeVec, strings.NewReader(header+expected))
    57  	assert.NoError(t, err)
    58  
    59  	g.Sub(ctx, 1)
    60  	expected = `
    61          testscope:unittest{domain="dev",lp="",project="flyte",task="",wf=""} 41
    62  	`
    63  	err = testutil.CollectAndCompare(g.GaugeVec, strings.NewReader(header+expected))
    64  	assert.NoError(t, err)
    65  
    66  	g.SetToCurrentTime(ctx)
    67  }
    68  
    69  func TestWithAdditionalLabels(t *testing.T) {
    70  	UnsetMetricKeys()
    71  	assert.NotPanics(t, func() {
    72  		SetMetricKeys(contextutils.ProjectKey, contextutils.DomainKey, contextutils.WorkflowIDKey, contextutils.TaskIDKey, contextutils.LaunchPlanIDKey)
    73  	})
    74  
    75  	scope := promutils.NewScope("testscope")
    76  	ctx := context.Background()
    77  	ctx = contextutils.WithProjectDomain(ctx, "flyte", "dev")
    78  	g := NewGauge("unittestlabeled", "some desc", scope, AdditionalLabelsOption{Labels: []string{"bearing"}})
    79  	assert.NotNil(t, g)
    80  
    81  	const header = `
    82  		# HELP testscope:unittestlabeled some desc
    83  		# TYPE testscope:unittestlabeled gauge
    84  	`
    85  
    86  	g.Inc(ctx)
    87  	var expected = `
    88          testscope:unittestlabeled{bearing="", domain="dev",lp="",project="flyte",task="",wf=""} 1
    89  	`
    90  	err := testutil.CollectAndCompare(g.GaugeVec, strings.NewReader(header+expected))
    91  	assert.NoError(t, err)
    92  
    93  	bearingKey := contextutils.Key("bearing")
    94  	ctx = context.WithValue(ctx, bearingKey, "123")
    95  	g.Set(ctx, 42)
    96  	expected = `
    97  		testscope:unittestlabeled{bearing="", domain="dev",lp="",project="flyte",task="",wf=""} 1
    98  		testscope:unittestlabeled{bearing="123", domain="dev",lp="",project="flyte",task="",wf=""} 42
    99  	`
   100  	err = testutil.CollectAndCompare(g.GaugeVec, strings.NewReader(header+expected))
   101  	assert.NoError(t, err)
   102  
   103  	g.Add(ctx, 1)
   104  	expected = `
   105  		testscope:unittestlabeled{bearing="", domain="dev",lp="",project="flyte",task="",wf=""} 1
   106  		testscope:unittestlabeled{bearing="123", domain="dev",lp="",project="flyte",task="",wf=""} 43
   107  	`
   108  	err = testutil.CollectAndCompare(g.GaugeVec, strings.NewReader(header+expected))
   109  	assert.NoError(t, err)
   110  
   111  	g.Dec(ctx)
   112  	expected = `
   113  		testscope:unittestlabeled{bearing="", domain="dev",lp="",project="flyte",task="",wf=""} 1
   114  		testscope:unittestlabeled{bearing="123", domain="dev",lp="",project="flyte",task="",wf=""} 42
   115  	`
   116  	err = testutil.CollectAndCompare(g.GaugeVec, strings.NewReader(header+expected))
   117  	assert.NoError(t, err)
   118  
   119  	g.Sub(ctx, 42)
   120  	expected = `
   121  		testscope:unittestlabeled{bearing="", domain="dev",lp="",project="flyte",task="",wf=""} 1
   122  		testscope:unittestlabeled{bearing="123", domain="dev",lp="",project="flyte",task="",wf=""} 0
   123  	`
   124  	err = testutil.CollectAndCompare(g.GaugeVec, strings.NewReader(header+expected))
   125  	assert.NoError(t, err)
   126  }