github.com/pyroscope-io/pyroscope@v0.37.3-0.20230725203016-5f6947968bd0/pkg/exporter/exporter_test.go (about)

     1  package exporter
     2  
     3  import (
     4  	"testing"
     5  
     6  	"github.com/prometheus/client_golang/prometheus"
     7  	"github.com/prometheus/client_golang/prometheus/testutil"
     8  
     9  	"github.com/pyroscope-io/pyroscope/pkg/agent/spy"
    10  	"github.com/pyroscope-io/pyroscope/pkg/config"
    11  	"github.com/pyroscope-io/pyroscope/pkg/storage"
    12  	"github.com/pyroscope-io/pyroscope/pkg/storage/segment"
    13  	"github.com/pyroscope-io/pyroscope/pkg/storage/tree"
    14  )
    15  
    16  func TestObserve(t *testing.T) {
    17  	rules := config.MetricsExportRules{
    18  		"app_name_cpu_total": {
    19  			Expr: `app.name.cpu`,
    20  		},
    21  		"app_name_cpu_total_foo": {
    22  			Expr:    `app.name.cpu{foo="bar"}`,
    23  			Node:    "total",
    24  			GroupBy: []string{"foo"},
    25  		},
    26  		"app_name_cpu_abc": {
    27  			Expr:    `app.name.cpu{foo=~"b.*"}`,
    28  			Node:    "^a;b;c$",
    29  			GroupBy: []string{"foo"},
    30  		},
    31  		"app_name_cpu_ab": {
    32  			Expr:    `app.name.cpu{foo=~"b.*"}`,
    33  			Node:    "a;b",
    34  			GroupBy: []string{"foo"},
    35  		},
    36  		"another_app": {
    37  			Expr:    `another.app.cpu`,
    38  			GroupBy: []string{"foo"},
    39  		},
    40  	}
    41  
    42  	exporter, _ := NewExporter(rules, prometheus.NewRegistry())
    43  	k := observe(exporter, "app.name.cpu{foo=bar,bar=baz}")
    44  
    45  	requireRuleCounterValue(t, exporter, "app_name_cpu_total", k, 5)
    46  	requireRuleCounterValue(t, exporter, "app_name_cpu_total_foo", k, 5)
    47  	requireRuleCounterValue(t, exporter, "app_name_cpu_abc", k, 2)
    48  	requireRuleCounterValue(t, exporter, "app_name_cpu_ab", k, 3)
    49  	requireNoCounter(t, exporter, "another_app", k)
    50  }
    51  
    52  const testRuleName = "app_name_cpu_total"
    53  
    54  func TestObserveNoMatch(t *testing.T) {
    55  	rules := config.MetricsExportRules{
    56  		testRuleName: {
    57  			Expr: `app.name.cpu`,
    58  		},
    59  	}
    60  
    61  	exporter, _ := NewExporter(rules, prometheus.NewRegistry())
    62  	k := observe(exporter, "another.app.name.cpu{foo=bar,bar=baz}")
    63  
    64  	requireNoCounter(t, exporter, testRuleName, k)
    65  }
    66  
    67  func TestGroupBy(t *testing.T) {
    68  	rules := config.MetricsExportRules{
    69  		testRuleName: {
    70  			Expr:    `app.name.cpu`,
    71  			GroupBy: []string{"foo"},
    72  		},
    73  	}
    74  
    75  	exporter, _ := NewExporter(rules, prometheus.NewRegistry())
    76  	k1 := observe(exporter, `app.name.cpu{foo=bar_a,bar=a}`)
    77  	k2 := observe(exporter, `app.name.cpu{foo=bar_a,bar=b}`)
    78  	k3 := observe(exporter, `app.name.cpu{foo=bar_b,bar=c}`)
    79  	k4 := observe(exporter, `app.name.cpu{}`)
    80  
    81  	requireRuleCounterValue(t, exporter, testRuleName, k1, 10)
    82  	requireRuleCounterValue(t, exporter, testRuleName, k2, 10)
    83  	requireRuleCounterValue(t, exporter, testRuleName, k3, 5)
    84  	requireRuleCounterValue(t, exporter, testRuleName, k4, 5)
    85  }
    86  
    87  func TestNoGroupBy(t *testing.T) {
    88  	rules := config.MetricsExportRules{
    89  		testRuleName: {
    90  			Expr: `app.name.cpu`,
    91  		},
    92  	}
    93  
    94  	exporter, _ := NewExporter(rules, prometheus.NewRegistry())
    95  	k1 := observe(exporter, "app.name.cpu{foo=bar}")
    96  	k2 := observe(exporter, "app.name.cpu{foo=baz}")
    97  
    98  	requireRuleCounterValue(t, exporter, testRuleName, k1, 10)
    99  	requireRuleCounterValue(t, exporter, testRuleName, k2, 10)
   100  }
   101  
   102  func TestSampleUnits(t *testing.T) {
   103  	rules := config.MetricsExportRules{
   104  		testRuleName: {Expr: "app.name.cpu"},
   105  	}
   106  
   107  	exporter, _ := NewExporter(rules, prometheus.NewRegistry())
   108  	k, _ := segment.ParseKey("app.name.cpu")
   109  	i := &storage.PutInput{Key: k, SampleRate: 100, Units: spy.ProfileCPU.Units()}
   110  	o, _ := exporter.Evaluate(i)
   111  	createTree().Iterate(observeCallback(o))
   112  
   113  	requireRuleCounterValue(t, exporter, testRuleName, k, 0.05)
   114  }
   115  
   116  func getCounter(e *MetricsExporter, name string, k *segment.Key) prometheus.Counter {
   117  	r, ok := e.rules[name]
   118  	if !ok {
   119  		return nil
   120  	}
   121  	return r.ctr.With(r.promLabels(k))
   122  }
   123  
   124  func requireNoCounter(t *testing.T, e *MetricsExporter, name string, k *segment.Key) {
   125  	r, ok := e.rules[name]
   126  	if !ok || r.ctr.Delete(r.promLabels(k)) {
   127  		t.Fatalf("Unexpected counter %s (%v)", name, k)
   128  	}
   129  }
   130  
   131  func requireRuleCounterValue(t *testing.T, e *MetricsExporter, name string, k *segment.Key, v float64) {
   132  	if actual := testutil.ToFloat64(getCounter(e, name, k)); v != actual {
   133  		t.Fatalf("Expected value %v got %v; counter %s (%v)", v, actual, name, k)
   134  	}
   135  }
   136  
   137  func observe(e *MetricsExporter, key string) *segment.Key {
   138  	k, _ := segment.ParseKey(key)
   139  	i := &storage.PutInput{Key: k, Units: spy.ProfileInuseObjects.Units()}
   140  	if o, ok := e.Evaluate(i); ok {
   141  		createTree().Iterate(observeCallback(o))
   142  	}
   143  	return k
   144  }
   145  
   146  func createTree() *tree.Tree {
   147  	t := tree.New()
   148  	t.Insert([]byte("a;b"), uint64(1))
   149  	t.Insert([]byte("a;c"), uint64(2))
   150  	t.Insert([]byte("a;b;c"), uint64(2))
   151  	return t
   152  }
   153  
   154  func observeCallback(o storage.SampleObserver) func([]byte, uint64) {
   155  	return func(key []byte, val uint64) {
   156  		// Key has ;; prefix.
   157  		if len(key) > 2 && val != 0 {
   158  			o.Observe(key[2:], int(val))
   159  		}
   160  	}
   161  }