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

     1  package middleware
     2  
     3  import (
     4  	"errors"
     5  	"net"
     6  	"testing"
     7  
     8  	"github.com/prometheus/client_golang/prometheus"
     9  	"github.com/prometheus/client_golang/prometheus/testutil"
    10  	"github.com/stretchr/testify/assert"
    11  )
    12  
    13  type fakeListener struct {
    14  	net.Listener
    15  	acceptErr error
    16  	closeErr  error
    17  }
    18  
    19  type fakeConn struct {
    20  	net.Conn
    21  	closeErr error
    22  }
    23  
    24  func (c *fakeConn) Close() error {
    25  	return c.closeErr
    26  }
    27  
    28  func (c *fakeListener) Accept() (net.Conn, error) {
    29  	return &fakeConn{closeErr: c.closeErr}, c.acceptErr
    30  }
    31  
    32  func TestCountingListener(t *testing.T) {
    33  	g := prometheus.NewGauge(prometheus.GaugeOpts{
    34  		Namespace: "test",
    35  		Name:      "gauge",
    36  	})
    37  
    38  	fake := &fakeListener{}
    39  	l := CountingListener(fake, g)
    40  	assert.Equal(t, float64(0), testutil.ToFloat64(g))
    41  
    42  	// Accepting connections should increment the gauge.
    43  	c1, err := l.Accept()
    44  	assert.NoError(t, err)
    45  	assert.Equal(t, float64(1), testutil.ToFloat64(g))
    46  	c2, err := l.Accept()
    47  	assert.NoError(t, err)
    48  	assert.Equal(t, float64(2), testutil.ToFloat64(g))
    49  
    50  	// Closing connections should decrement the gauge.
    51  	assert.NoError(t, c1.Close())
    52  	assert.Equal(t, float64(1), testutil.ToFloat64(g))
    53  	assert.NoError(t, c2.Close())
    54  	assert.Equal(t, float64(0), testutil.ToFloat64(g))
    55  
    56  	// Duplicate calls to Close should not decrement.
    57  	assert.NoError(t, c1.Close())
    58  	assert.Equal(t, float64(0), testutil.ToFloat64(g))
    59  
    60  	// Accept errors should not cause an increment.
    61  	fake.acceptErr = errors.New("accept")
    62  	_, err = l.Accept()
    63  	assert.Error(t, err)
    64  	assert.Equal(t, float64(0), testutil.ToFloat64(g))
    65  
    66  	// Close errors should still decrement.
    67  	fake.acceptErr = nil
    68  	fake.closeErr = errors.New("close")
    69  	c3, err := l.Accept()
    70  	assert.NoError(t, err)
    71  	assert.Equal(t, float64(1), testutil.ToFloat64(g))
    72  	assert.Error(t, c3.Close())
    73  	assert.Equal(t, float64(0), testutil.ToFloat64(g))
    74  }