1package test
2
3import (
4	"context"
5	"fmt"
6	"sort"
7	"strings"
8	"time"
9
10	"go.opencensus.io/metric/metricdata"
11	"go.opencensus.io/metric/metricexport"
12	"go.opencensus.io/stats/view"
13)
14
15// Exporter keeps exported metric data in memory to aid in testing the instrumentation.
16//
17// Metrics can be retrieved with `GetPoint()`. In order to deterministically retrieve the most recent values, you must first invoke `ReadAndExport()`.
18type Exporter struct {
19	// points is a map from a label signature to the latest value for the time series represented by the signature.
20	// Use function `labelSignature` to get a signature from a `metricdata.Metric`.
21	points       map[string]metricdata.Point
22	metricReader *metricexport.Reader
23}
24
25var _ metricexport.Exporter = &Exporter{}
26
27// NewExporter returns a new exporter.
28func NewExporter(metricReader *metricexport.Reader) *Exporter {
29	return &Exporter{points: make(map[string]metricdata.Point), metricReader: metricReader}
30}
31
32// ExportMetrics records the view data.
33func (e *Exporter) ExportMetrics(ctx context.Context, data []*metricdata.Metric) error {
34	for _, metric := range data {
35		for _, ts := range metric.TimeSeries {
36			signature := labelSignature(metric.Descriptor.Name, labelObjectsToKeyValue(metric.Descriptor.LabelKeys, ts.LabelValues))
37			e.points[signature] = ts.Points[len(ts.Points)-1]
38		}
39	}
40	return nil
41}
42
43// GetPoint returns the latest point for the time series identified by the given labels.
44func (e *Exporter) GetPoint(metricName string, labels map[string]string) (metricdata.Point, bool) {
45	v, ok := e.points[labelSignature(metricName, labelMapToKeyValue(labels))]
46	return v, ok
47}
48
49// ReadAndExport reads the current values for all metrics and makes them available to this exporter.
50func (e *Exporter) ReadAndExport() {
51	// The next line forces the view worker to process all stats.Record* calls that
52	// happened within Store() before the call to ReadAndExport below. This abuses the
53	// worker implementation to work around lack of synchronization.
54	// TODO(jkohen,rghetia): figure out a clean way to make this deterministic.
55	view.SetReportingPeriod(time.Minute)
56	e.metricReader.ReadAndExport(e)
57}
58
59// String defines the ``native'' format for the exporter.
60func (e *Exporter) String() string {
61	return fmt.Sprintf("points{%v}", e.points)
62}
63
64type keyValue struct {
65	Key   string
66	Value string
67}
68
69func sortKeyValue(kv []keyValue) {
70	sort.Slice(kv, func(i, j int) bool { return kv[i].Key < kv[j].Key })
71}
72
73func labelMapToKeyValue(labels map[string]string) []keyValue {
74	kv := make([]keyValue, 0, len(labels))
75	for k, v := range labels {
76		kv = append(kv, keyValue{Key: k, Value: v})
77	}
78	sortKeyValue(kv)
79	return kv
80}
81
82func labelObjectsToKeyValue(keys []metricdata.LabelKey, values []metricdata.LabelValue) []keyValue {
83	if len(keys) != len(values) {
84		panic("keys and values must have the same length")
85	}
86	kv := make([]keyValue, 0, len(values))
87	for i := range keys {
88		if values[i].Present {
89			kv = append(kv, keyValue{Key: keys[i].Key, Value: values[i].Value})
90		}
91	}
92	sortKeyValue(kv)
93	return kv
94}
95
96// labelSignature returns a string that uniquely identifies the list of labels given in the input.
97func labelSignature(metricName string, kv []keyValue) string {
98	var builder strings.Builder
99	for _, x := range kv {
100		builder.WriteString(x.Key)
101		builder.WriteString(x.Value)
102	}
103	return fmt.Sprintf("%s{%s}", metricName, builder.String())
104}
105