1// Copyright 2017 The Go Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style
3// license that can be found in the LICENSE file.
4
5package pprof
6
7import (
8	"context"
9	"fmt"
10	"reflect"
11	"testing"
12)
13
14func TestSetGoroutineLabels(t *testing.T) {
15	sync := make(chan struct{})
16
17	wantLabels := map[string]string{}
18	if gotLabels := getProfLabel(); !reflect.DeepEqual(gotLabels, wantLabels) {
19		t.Errorf("Expected parent goroutine's profile labels to be empty before test, got %v", gotLabels)
20	}
21	go func() {
22		if gotLabels := getProfLabel(); !reflect.DeepEqual(gotLabels, wantLabels) {
23			t.Errorf("Expected child goroutine's profile labels to be empty before test, got %v", gotLabels)
24		}
25		sync <- struct{}{}
26	}()
27	<-sync
28
29	wantLabels = map[string]string{"key": "value"}
30	ctx := WithLabels(context.Background(), Labels("key", "value"))
31	SetGoroutineLabels(ctx)
32	if gotLabels := getProfLabel(); !reflect.DeepEqual(gotLabels, wantLabels) {
33		t.Errorf("parent goroutine's profile labels: got %v, want %v", gotLabels, wantLabels)
34	}
35	go func() {
36		if gotLabels := getProfLabel(); !reflect.DeepEqual(gotLabels, wantLabels) {
37			t.Errorf("child goroutine's profile labels: got %v, want %v", gotLabels, wantLabels)
38		}
39		sync <- struct{}{}
40	}()
41	<-sync
42
43	wantLabels = map[string]string{}
44	ctx = context.Background()
45	SetGoroutineLabels(ctx)
46	if gotLabels := getProfLabel(); !reflect.DeepEqual(gotLabels, wantLabels) {
47		t.Errorf("Expected parent goroutine's profile labels to be empty, got %v", gotLabels)
48	}
49	go func() {
50		if gotLabels := getProfLabel(); !reflect.DeepEqual(gotLabels, wantLabels) {
51			t.Errorf("Expected child goroutine's profile labels to be empty, got %v", gotLabels)
52		}
53		sync <- struct{}{}
54	}()
55	<-sync
56}
57
58func TestDo(t *testing.T) {
59	wantLabels := map[string]string{}
60	if gotLabels := getProfLabel(); !reflect.DeepEqual(gotLabels, wantLabels) {
61		t.Errorf("Expected parent goroutine's profile labels to be empty before Do, got %v", gotLabels)
62	}
63
64	Do(context.Background(), Labels("key1", "value1", "key2", "value2"), func(ctx context.Context) {
65		wantLabels := map[string]string{"key1": "value1", "key2": "value2"}
66		if gotLabels := getProfLabel(); !reflect.DeepEqual(gotLabels, wantLabels) {
67			t.Errorf("parent goroutine's profile labels: got %v, want %v", gotLabels, wantLabels)
68		}
69
70		sync := make(chan struct{})
71		go func() {
72			wantLabels := map[string]string{"key1": "value1", "key2": "value2"}
73			if gotLabels := getProfLabel(); !reflect.DeepEqual(gotLabels, wantLabels) {
74				t.Errorf("child goroutine's profile labels: got %v, want %v", gotLabels, wantLabels)
75			}
76			sync <- struct{}{}
77		}()
78		<-sync
79
80	})
81
82	wantLabels = map[string]string{}
83	if gotLabels := getProfLabel(); !reflect.DeepEqual(gotLabels, wantLabels) {
84		fmt.Printf("%#v", gotLabels)
85		fmt.Printf("%#v", wantLabels)
86		t.Errorf("Expected parent goroutine's profile labels to be empty after Do, got %v", gotLabels)
87	}
88}
89
90func getProfLabel() map[string]string {
91	l := (*labelMap)(runtime_getProfLabel())
92	if l == nil {
93		return map[string]string{}
94	}
95	return *l
96}
97