github.com/weaveworks/common@v0.0.0-20230728070032-dd9e68f319d5/instrument/instrument_test.go (about)

     1  package instrument_test
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/prometheus/client_golang/prometheus"
    10  	"github.com/stretchr/testify/assert"
    11  
    12  	"github.com/weaveworks/common/instrument"
    13  )
    14  
    15  func TestNewHistogramCollector(t *testing.T) {
    16  	m := prometheus.NewHistogramVec(prometheus.HistogramOpts{
    17  		Namespace: "test",
    18  		Subsystem: "instrumentation",
    19  		Name:      "foo",
    20  		Help:      "",
    21  		Buckets:   prometheus.DefBuckets,
    22  	}, instrument.HistogramCollectorBuckets)
    23  	c := instrument.NewHistogramCollector(m)
    24  	assert.NotNil(t, c)
    25  }
    26  
    27  type spyCollector struct {
    28  	before    bool
    29  	after     bool
    30  	afterCode string
    31  }
    32  
    33  func (c *spyCollector) Register() {
    34  }
    35  
    36  // Before collects for the upcoming request.
    37  func (c *spyCollector) Before(ctx context.Context, method string, start time.Time) {
    38  	c.before = true
    39  }
    40  
    41  // After collects when the request is done.
    42  func (c *spyCollector) After(ctx context.Context, method, statusCode string, start time.Time) {
    43  	c.after = true
    44  	c.afterCode = statusCode
    45  }
    46  
    47  func TestCollectedRequest(t *testing.T) {
    48  	c := &spyCollector{}
    49  	fcalled := false
    50  	instrument.CollectedRequest(context.Background(), "test", c, nil, func(_ context.Context) error {
    51  		fcalled = true
    52  		return nil
    53  	})
    54  	assert.True(t, fcalled)
    55  	assert.True(t, c.before)
    56  	assert.True(t, c.after)
    57  	assert.Equal(t, "200", c.afterCode)
    58  }
    59  
    60  func TestCollectedRequest_Error(t *testing.T) {
    61  	c := &spyCollector{}
    62  	instrument.CollectedRequest(context.Background(), "test", c, nil, func(_ context.Context) error {
    63  		return errors.New("boom")
    64  	})
    65  	assert.True(t, c.before)
    66  	assert.True(t, c.after)
    67  	assert.Equal(t, "500", c.afterCode)
    68  }