github.com/prysmaticlabs/prysm@v1.4.4/beacon-chain/powchain/prometheus_test.go (about)

     1  package powchain
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/prometheus/client_golang/prometheus"
     9  	"github.com/prysmaticlabs/prysm/shared/testutil/assert"
    10  )
    11  
    12  // TestCleanup ensures that the cleanup function unregisters the prometheus.Collection
    13  // also tests the interchangability of the explicit prometheus Register/Unregister
    14  // and the implicit methods within the collector implementation
    15  func TestCleanup(t *testing.T) {
    16  	ctx := context.Background()
    17  	pc, err := NewPowchainCollector(ctx)
    18  	assert.NoError(t, err, "Uxpected error caling NewPowchainCollector")
    19  	unregistered := pc.unregister()
    20  	assert.Equal(t, true, unregistered, "PowchainCollector.unregister did not return true (via prometheus.DefaultRegistry)")
    21  	// PowchainCollector is a prometheus.Collector, so we should be able to register it again
    22  	err = prometheus.Register(pc)
    23  	assert.NoError(t, err, "Got error from prometheus.Register after unregistering PowchainCollector")
    24  	// even if it somehow gets registered somewhere else, unregister should work
    25  	unregistered = pc.unregister()
    26  	assert.Equal(t, true, unregistered, "PowchainCollector.unregister failed on the second attempt")
    27  	// and so we should be able to register it again
    28  	err = prometheus.Register(pc)
    29  	assert.NoError(t, err, "Got error from prometheus.Register on the second attempt")
    30  	// ok clean it up one last time for real :)
    31  	unregistered = prometheus.Unregister(pc)
    32  	assert.Equal(t, true, unregistered, "prometheus.Unregister failed to unregister PowchainCollector on final cleanup")
    33  }
    34  
    35  // TestCancelation tests that canceling the context passed into
    36  // NewPowchainCollector cleans everything up as expected. This
    37  // does come at the cost of an extra channel cluttering up
    38  // PowchainCollector, just for this test.
    39  func TestCancelation(t *testing.T) {
    40  	ctx, cancel := context.WithCancel(context.Background())
    41  	pc, err := NewPowchainCollector(ctx)
    42  	assert.NoError(t, err, "Uxpected error caling NewPowchainCollector")
    43  	ticker := time.NewTicker(10 * time.Second)
    44  	cancel()
    45  	select {
    46  	case <-ticker.C:
    47  		t.Error("Hit timeout waiting for cancel() to cleanup PowchainCollector")
    48  	case <-pc.finishChan:
    49  		break
    50  	}
    51  	err = prometheus.Register(pc)
    52  	assert.NoError(t, err, "Got error from prometheus.Register after unregistering PowchainCollector through canceled context")
    53  }