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