Repository: grpc/grpc-go
File: internal/testutils/stats/test_metrics_recorder.go (Go)

Code before:
	"sync"
	"testing"

	"github.com/google/go-cmp/cmp"
	estats "google.golang.org/grpc/experimental/stats"
	"google.golang.org/grpc/internal/testutils"
	"google.golang.org/grpc/stats"
)

// TestMetricsRecorder is a MetricsRecorder to be used in tests. It sends
// recording events on channels and provides helpers to check if certain events
// have taken place. It also persists metrics data keyed on the metrics
// descriptor.
type TestMetricsRecorder struct {
	t *testing.T

	intCountCh   *testutils.Channel
	floatCountCh *testutils.Channel
	intHistoCh   *testutils.Channel
	floatHistoCh *testutils.Channel
	intGaugeCh   *testutils.Channel

	// Mu protects Data.
	Mu sync.Mutex
	// Data is the most recent update for each metric name.
	Data map[estats.Metric]float64
}

func NewTestMetricsRecorder() *TestMetricsRecorder {
	return &TestMetricsRecorder{
		intCountCh:   testutils.NewChannelWithSize(10),
		floatCountCh: testutils.NewChannelWithSize(10),
		intHistoCh:   testutils.NewChannelWithSize(10),
		floatHistoCh: testutils.NewChannelWithSize(10),
		intGaugeCh:   testutils.NewChannelWithSize(10),

		Data: make(map[estats.Metric]float64),
	}
}

// ClearMetrics clears the metrics data stores of the test metrics recorder by
// setting all the data to 0.
func (r *TestMetricsRecorder) ClearMetrics() {
	r.Mu.Lock()
	defer r.Mu.Unlock()
	r.Data = make(map[estats.Metric]float64)
}

type MetricsData struct {
	Handle *estats.MetricDescriptor


Review comment (suggestion):
A getter may be better than exporting?

```go
func (...) Metric(name string) float64
```

Code after:
	"sync"

	"github.com/google/go-cmp/cmp"
	estats "google.golang.org/grpc/experimental/stats"
	"google.golang.org/grpc/internal/testutils"
	"google.golang.org/grpc/stats"
)

// TestMetricsRecorder is a MetricsRecorder to be used in tests. It sends
// recording events on channels and provides helpers to check if certain events
// have taken place. It also persists metrics data keyed on the metrics
// descriptor.
type TestMetricsRecorder struct {
	intCountCh   *testutils.Channel
	floatCountCh *testutils.Channel
	intHistoCh   *testutils.Channel
	floatHistoCh *testutils.Channel
	intGaugeCh   *testutils.Channel

	// mu protects data.
	mu sync.Mutex
	// data is the most recent update for each metric name.
	data map[estats.Metric]float64
}

// NewTestMetricsRecorder returns a new TestMetricsRecorder.
func NewTestMetricsRecorder() *TestMetricsRecorder {
	return &TestMetricsRecorder{
		intCountCh:   testutils.NewChannelWithSize(10),
		floatCountCh: testutils.NewChannelWithSize(10),
		intHistoCh:   testutils.NewChannelWithSize(10),
		floatHistoCh: testutils.NewChannelWithSize(10),
		intGaugeCh:   testutils.NewChannelWithSize(10),

		data: make(map[estats.Metric]float64),
	}
}

// Metric returns the most recent data for a metric, and whether this recorder
// has received data for a metric.
func (r *TestMetricsRecorder) Metric(name string) (float64, bool) {
	data, ok := r.data[estats.Metric(name)]
	return data, ok
}

// ClearMetrics clears the metrics data store of the test metrics recorder.
func (r *TestMetricsRecorder) ClearMetrics() {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.data = make(map[estats.Metric]float64)
}