1// Copyright 2018 Google LLC
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15package testutil
16
17import (
18	"log"
19	"sync"
20	"time"
21
22	"go.opencensus.io/plugin/ocgrpc"
23	"go.opencensus.io/stats/view"
24	"go.opencensus.io/trace"
25)
26
27// TestExporter is a test utility exporter. It should be created with NewtestExporter.
28type TestExporter struct {
29	mu    sync.Mutex
30	Spans []*trace.SpanData
31
32	Stats chan *view.Data
33	Views []*view.View
34}
35
36// NewTestExporter creates a TestExporter and registers it with OpenCensus.
37func NewTestExporter(views ...*view.View) *TestExporter {
38	if len(views) == 0 {
39		views = ocgrpc.DefaultClientViews
40	}
41	te := &TestExporter{Stats: make(chan *view.Data), Views: views}
42
43	view.RegisterExporter(te)
44	view.SetReportingPeriod(time.Millisecond)
45	if err := view.Register(views...); err != nil {
46		log.Fatal(err)
47	}
48
49	trace.RegisterExporter(te)
50	trace.ApplyConfig(trace.Config{DefaultSampler: trace.AlwaysSample()})
51
52	return te
53}
54
55// ExportSpan exports a span.
56func (te *TestExporter) ExportSpan(s *trace.SpanData) {
57	te.mu.Lock()
58	defer te.mu.Unlock()
59	te.Spans = append(te.Spans, s)
60}
61
62// ExportView exports a view.
63func (te *TestExporter) ExportView(vd *view.Data) {
64	if len(vd.Rows) > 0 {
65		select {
66		case te.Stats <- vd:
67		default:
68		}
69	}
70}
71
72// Unregister unregisters the exporter from OpenCensus.
73func (te *TestExporter) Unregister() {
74	view.Unregister(te.Views...)
75	view.UnregisterExporter(te)
76	trace.UnregisterExporter(te)
77	view.SetReportingPeriod(0) // reset to default value
78}
79