Repository: census-instrumentation/opencensus-go
Pull request #231 (2017-12-11 15:03:29)

Update SetDefaultSampler to revert to default sampler
Hey there!  I was looking through the issues for one that would get me into the code and found #197 

I'm not 100% sure this is what is needed, but, it's fun to get in and see how it is built!



rakyll: Thanks for the PR! We haven't agreed this should be the expected behavior yet but @jcd2 can provide feedback.

Patch:
diff --git a/README.md b/README.md
index b6a4246..ea679d6 100644
--- a/README.md
+++ b/README.md
@@ -136,14 +136,18 @@ if err := stats.DeleteMeasure(m); err != nil {
 
 ### Creating an aggregation
 
-Currently only 2 types of aggregations are supported. The CountAggregation is used to count
+Currently 4 types of aggregations are supported. The CountAggregation is used to count
 the number of times a sample was recorded. The DistributionAggregation is used to
-provide a histogram of the values of the samples.
+provide a histogram of the values of the samples. The SumAggregation is used to
+sum up all sample values. The MeanAggregation is used to calculate the mean of
+sample values.
 
 [embedmd]:# (stats.go aggs)
 ```go
 distAgg := stats.DistributionAggregation([]float64{0, 1 << 32, 2 << 32, 3 << 32})
 countAgg := stats.CountAggregation{}
+sumAgg := stats.SumAggregation{}
+meanAgg := stats.MeanAggregation{}
 ```
 
 ### Create an aggregation window
@@ -170,7 +174,7 @@ Create and register a view:
 
 [embedmd]:# (stats.go view)
 ```go
-view := stats.NewView(
+view, err := stats.NewView(
 	"my.org/video_size_distribution",
 	"distribution of processed video size over time",
 	nil,
@@ -178,6 +182,9 @@ view := stats.NewView(
 	distAgg,
 	cum,
 )
+if err != nil {
+	log.Fatalf("cannot create view: %v", err)
+}
 if err := stats.RegisterView(view); err != nil {
 	log.Fatal(err)
 }
diff --git a/exporter/stats/prometheus/prometheus.go b/exporter/stats/prometheus/prometheus.go
index a92af36..ffab6c2 100644
--- a/exporter/stats/prometheus/prometheus.go
+++ b/exporter/stats/prometheus/prometheus.go
@@ -72,9 +72,20 @@ func NewExporter(o Options) (*Exporter, error) {
 var _ http.Handler = (*Exporter)(nil)
 var _ stats.Exporter = (*Exporter)(nil)
 
+func allowedWindowType(v *stats.View) bool {
+	// TODO: (@rakyll, @odeke-em): Only the cumulative window will
+	// be exported in this version. Support others in the future.
+	// See Issue https://github.com/census-instrumentation/opencensus-go/issues/214
+	_, ok := v.Window().(stats.Cumulative)
+	return ok
+}
+
 func (c *collector) registerViews(views ...*stats.View) {
 	count := 0
 	for _, view := range views {
+		if !allowedWindowType(view) {
+			continue
+		}
 		sig := viewSignature(c.opts.Namespace, view)
 		c.registeredViewsMu.Lock()
 		_, ok := c.registeredViews[sig]
@@ -170,14 +181,11 @@ func (c *collector) Describe(ch chan<- *prometheus.Desc) {
 // Collect is invoked everytime a prometheus.Gatherer is run
 // for example when the HTTP endpoint is invoked by Prometheus.
 func (c *collector) Collect(ch chan<- prometheus.Metric) {
-	c.mu.Lock()
-	views := make(map[string]*stats.ViewData, len(c.viewData))
-	for i, vd := range c.viewData {
-		views[i] = vd
-	}
-	c.mu.Unlock()
-
 	for _, vd := range c.viewData {
+		if !allowedWindowType(vd.View) {
+			continue
+		}
+
 		sig := viewSignature(c.opts.Namespace, vd.View)
 		c.registeredViewsMu.Lock()
 		desc := c.registeredViews[sig]
diff --git a/exporter/stats/stackdriver/stackdriver.go b/exporter/stats/stackdriver/stackdriver.go
index 7689401..0dd4709 100644
--- a/exporter/stats/stackdriver/stackdriver.go
+++ b/exporter/stats/stackdriver/stackdriver.go
@@ -29,11 +29,9 @@ import (
 	"sync"
 	"time"
 
-	"go.opencensus.io/tag"
-
 	"go.opencensus.io/internal"
-
 	"go.opencensus.io/stats"
+	"go.opencensus.io/tag"
 
 	monitoring "cloud.google.com/go/monitoring/apiv3"
 	timestamp "github.com/golang/protobuf/ptypes/timestamp"
@@ -41,6 +39,7 @@ import (
 	"google.golang.org/api/support/bundler"
 	distributionpb "google.golang.org/genproto/googleapis/api/distribution"
 	labelpb "google.golang.org/genproto/googleapis/api/label"
+	"google.golang.org/genproto/googleapis/api/metric"
 	metricpb "google.golang.org/genproto/googleapis/api/metric"
 	monitoredrespb "google.golang.org/genproto/googleapis/api/monitoredres"
 	monitoringpb "google.golang.org/genproto/googleapis/monitoring/v3"
@@ -222,7 +221,7 @@ func (e *Exporter) createMeasure(ctx context.Context, vd *stats.ViewData) error
 	}
 
 	metricName := monitoring.MetricMetricDescriptorPath(e.o.ProjectID, namespacedViewName(viewName, true))
-	md, err := e.c.GetMetricDescriptor(ctx, &monitoringpb.GetMetricDescriptorRequest{
+	md, err := getMetricDescriptor(ctx, e.c, &monitoringpb.GetMetricDescriptorRequest{
 		Name: metricName,
 	})
 	if err == nil {
@@ -261,10 +260,10 @@ func (e *Exporter) createMeasure(ctx context.Context, vd *stats.ViewData) error
 		return fmt.Errorf("unsupported window type: %T", window)
 	}
 
-	md, err = e.c.CreateMetricDescriptor(ctx, &monitoringpb.CreateMetricDescriptorRequest{
+	md, err = createMetricDescriptor(ctx, e.c, &monitoringpb.CreateMetricDescriptorRequest{
 		Name: monitoring.MetricProjectPath(e.o.ProjectID),
 		MetricDescriptor: &metricpb.MetricDescriptor{
-			DisplayName: viewName,
+			DisplayName: path.Join("OpenCensus", viewName),
 			Description: m.Description(),
 			Unit:        m.Unit(),
 			Type:        namespacedViewName(viewName, false),
@@ -429,3 +428,11 @@ func equalAggWindowTagKeys(md *metricpb.MetricDescriptor, agg stats.Aggregation,
 
 	return nil
 }
+
+var createMetricDescriptor = func(ctx context.Context, c *monitoring.MetricClient, mdr *monitoringpb.CreateMetricDescriptorRequest) (*metric.MetricDescriptor, error) {
+	return c.CreateMetricDescriptor(ctx, mdr)
+}
+
+var getMetricDescriptor = func(ctx context.Context, c *monitoring.MetricClient, mdr *monitoringpb.GetMetricDescriptorRequest) (*metric.MetricDescriptor, error) {
+	return c.GetMetricDescriptor(ctx, mdr)
+}
diff --git a/internal/readme/source.md b/internal/readme/source.md
index bf4780c..87c6f07 100644
--- a/internal/readme/source.md
+++ b/internal/readme/source.md
@@ -80,9 +80,11 @@ another measure with the same name):
 
 ### Creating an aggregation
 
-Currently only 2 types of aggregations are supported. The CountAggregation is used to count
+Currently 4 types of aggregations are supported. The CountAggregation is used to count
 the number of times a sample was recorded. The DistributionAggregation is used to
-provide a histogram of the values of the samples.
+provide a histogram of the values of the samples. The SumAggregation is used to
+sum up all sample values. The MeanAggregation is used to calculate the mean of
+sample values.
 
 [embedmd]:# (stats.go aggs)
 
diff --git a/internal/readme/stats.go b/internal/readme/stats.go
index 7bfad5a..631716f 100644
--- a/internal/readme/stats.go
+++ b/internal/readme/stats.go
@@ -55,9 +55,11 @@ func statsExamples() {
 	// START aggs
 	distAgg := stats.DistributionAggregation([]float64{0, 1 << 32, 2 << 32, 3 << 32})
 	countAgg := stats.CountAggregation{}
+	sumAgg := stats.SumAggregation{}
+	meanAgg := stats.MeanAggregation{}
 	// END aggs
 
-	_, _ = distAgg, countAgg
+	_, _, _, _ = distAgg, countAgg, sumAgg, meanAgg
 
 	// START windows
 	interval := stats.Interval{
diff --git a/tag/map_codec.go b/tag/map_codec.go
index 49cc433..afe699c 100644
--- a/tag/map_codec.go
+++ b/tag/map_codec.go
@@ -189,7 +189,7 @@ func Decode(bytes []byte) (*Map, error) {
 
 	version := eg.readByte()
 	if version > tagsVersionID {
-		return nil, fmt.Errorf("decode failed; unsupported version: %q; supports only up to: %q", version, tagsVersionID)
+		return nil, fmt.Errorf("cannot decode: unsupported version: %q; supports only up to: %q", version, tagsVersionID)
 	}
 
 	for !eg.readEnded() {
@@ -199,7 +199,7 @@ func Decode(bytes []byte) (*Map, error) {
 		case keyTypeString:
 			break
 		default:
-			return nil, fmt.Errorf("decode failed; invalid key type: %q", typ)
+			return nil, fmt.Errorf("cannot decode: invalid key type: %q", typ)
 		}
 
 		k, err := eg.readBytesWithVarintLen()
@@ -214,11 +214,13 @@ func Decode(bytes []byte) (*Map, error) {
 
 		key, err := NewKey(string(k))
 		if err != nil {
-			// TODO(acetechnologist): log that key received on the wire and its value was ignored
-			continue
+			return nil, err // no partial failures
 		}
-		ts.upsert(key, string(v))
+		val := string(v)
+		if !checkValue(val) {
+			return nil, errInvalid // no partial failures
+		}
+		ts.upsert(key, val)
 	}
-
 	return ts, nil
 }
diff --git a/trace/sampling.go b/trace/sampling.go
index 05bf295..8f177c9 100644
--- a/trace/sampling.go
+++ b/trace/sampling.go
@@ -23,13 +23,17 @@ const defaultSamplingProbability = 1e-4
 var defaultSampler Sampler
 
 func init() {
-	defaultSampler = ProbabilitySampler(defaultSamplingProbability)
+	defaultSampler = newDefaultSampler()
+}
+
+func newDefaultSampler() Sampler {
+	return ProbabilitySampler(defaultSamplingProbability)
 }
 
 // SetDefaultSampler sets the default sampler used when creating new spans.
 func SetDefaultSampler(sampler Sampler) {
 	if sampler == nil {
-		sampler = NeverSample()
+		sampler = newDefaultSampler()
 	}
 	mu.Lock()
 	defaultSampler = sampler


Test patch:
diff --git a/exporter/stats/prometheus/prometheus_test.go b/exporter/stats/prometheus/prometheus_test.go
new file mode 100644
index 0000000..f2141b4
--- /dev/null
+++ b/exporter/stats/prometheus/prometheus_test.go
@@ -0,0 +1,77 @@
+// Copyright 2017, OpenCensus Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package prometheus
+
+import (
+	"testing"
+
+	"go.opencensus.io/stats"
+
+	"github.com/prometheus/client_golang/prometheus"
+)
+
+func newView(agg stats.Aggregation, window stats.Window) *stats.View {
+	m, _ := stats.NewMeasureInt64("tests/foo1", "bytes", "byte")
+	view, _ := stats.NewView("foo", "bar", nil, m, agg, window)
+	return view
+}
+
+func TestOnlyCumulativeWindowSupported(t *testing.T) {
+	// See Issue https://github.com/census-instrumentation/opencensus-go/issues/214.
+	count1 := stats.CountData(1)
+	tests := []struct {
+		vds  *stats.ViewData
+		want int
+	}{
+		0: {
+			vds: &stats.ViewData{
+				View: newView(stats.CountAggregation{}, stats.Cumulative{}),
+			},
+			want: 0, // no rows present
+		},
+		1: {
+			vds: &stats.ViewData{
+				View: newView(stats.CountAggregation{}, stats.Cumulative{}),
+				Rows: []*stats.Row{
+					{nil, &count1},
+				},
+			},
+			want: 1,
+		},
+		2: {
+			vds: &stats.ViewData{
+				View: newView(stats.CountAggregation{}, stats.Interval{}),
+				Rows: []*stats.Row{
+					{nil, &count1},
+				},
+			},
+			want: 0,
+		},
+	}
+
+	for i, tt := range tests {
+		reg := prometheus.NewRegistry()
+		collector := newCollector(Options{}, reg)
+		collector.addViewData(tt.vds)
+		mm, err := reg.Gather()
+		if err != nil {
+			t.Errorf("#%d: Gather err: %v", i, err)
+		}
+		reg.Unregister(collector)
+		if got, want := len(mm), tt.want; got != want {
+			t.Errorf("#%d: got nil %v want nil %v", i, got, want)
+		}
+	}
+}
diff --git a/exporter/stats/stackdriver/stackdriver_test.go b/exporter/stats/stackdriver/stackdriver_test.go
index 4b49265..994c3ee 100644
--- a/exporter/stats/stackdriver/stackdriver_test.go
+++ b/exporter/stats/stackdriver/stackdriver_test.go
@@ -15,21 +15,23 @@
 package stackdriver
 
 import (
+	"context"
 	"reflect"
 	"testing"
 	"time"
 
-	distributionpb "google.golang.org/genproto/googleapis/api/distribution"
-	"google.golang.org/genproto/googleapis/api/label"
-
-	"go.opencensus.io/stats"
-	"go.opencensus.io/tag"
-
 	monitoring "cloud.google.com/go/monitoring/apiv3"
 	timestamp "github.com/golang/protobuf/ptypes/timestamp"
+	"go.opencensus.io/stats"
+	"go.opencensus.io/tag"
+	distributionpb "google.golang.org/genproto/googleapis/api/distribution"
+	"google.golang.org/genproto/googleapis/api/label"
+	"google.golang.org/genproto/googleapis/api/metric"
 	metricpb "google.golang.org/genproto/googleapis/api/metric"
 	monitoredrespb "google.golang.org/genproto/googleapis/api/monitoredres"
 	monitoringpb "google.golang.org/genproto/googleapis/monitoring/v3"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 )
 
 func TestExporter_makeReq(t *testing.T) {
@@ -525,6 +527,70 @@ func TestEqualAggWindowTagKeys(t *testing.T) {
 	}
 }
 
+func TestExporter_createMeasure(t *testing.T) {
+	oldGetMetricDescriptor := getMetricDescriptor
+	oldCreateMetricDescriptor := createMetricDescriptor
+
+	defer func() {
+		getMetricDescriptor = oldGetMetricDescriptor
+		createMetricDescriptor = oldCreateMetricDescriptor
+	}()
+
+	key, _ := tag.NewKey("test-key-one")
+	m, err := stats.NewMeasureFloat64("test-measure", "measure desc", "unit")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer stats.DeleteMeasure(m)
+
+	view, err := stats.NewView("cumview", "desc", []tag.Key{key}, m, stats.CountAggregation{}, stats.Cumulative{})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	data := stats.CountData(0)
+	vd := newTestCumViewData(view, time.Now(), time.Now(), &data, &data)
+
+	e := &Exporter{
+		createdViews: make(map[string]*metricpb.MetricDescriptor),
+	}
+
+	var getCalls, createCalls int
+	getMetricDescriptor = func(ctx context.Context, c *monitoring.MetricClient, mdr *monitoringpb.GetMetricDescriptorRequest) (*metric.MetricDescriptor, error) {
+		getCalls++
+		return nil, status.Error(codes.NotFound, "")
+	}
+	createMetricDescriptor = func(ctx context.Context, c *monitoring.MetricClient, mdr *monitoringpb.CreateMetricDescriptorRequest) (*metric.MetricDescriptor, error) {
+		createCalls++
+		return &metric.MetricDescriptor{
+			DisplayName: "display",
+			Description: "desc",
+			Unit:        "unit",
+			Type:        "hello",
+			MetricKind:  metricpb.MetricDescriptor_CUMULATIVE,
+			ValueType:   metricpb.MetricDescriptor_INT64,
+			Labels:      newLabelDescriptors(vd.View.TagKeys()),
+		}, nil
+	}
+
+	ctx := context.Background()
+	if err := e.createMeasure(ctx, vd); err != nil {
+		t.Errorf("Exporter.createMeasure() error = %v", err)
+	}
+	if err := e.createMeasure(ctx, vd); err != nil {
+		t.Errorf("Exporter.createMeasure() error = %v", err)
+	}
+	if count := getCalls; count != 1 {
+		t.Errorf("getMetricDescriptor needs to be called for once; called %v times", count)
+	}
+	if count := createCalls; count != 1 {
+		t.Errorf("createMetricDescriptor needs to be called for once; called %v times", count)
+	}
+	if count := len(e.createdViews); count != 1 {
+		t.Errorf("len(e.createdViews) = %v; want 1", count)
+	}
+}
+
 func newTestCumViewData(v *stats.View, start, end time.Time, data1, data2 stats.AggregationData) *stats.ViewData {
 	key, _ := tag.NewKey("test-key")
 	tag1 := tag.Tag{Key: key, Value: "test-value-1"}
diff --git a/tag/map_codec_test.go b/tag/map_codec_test.go
index f787ad7..fbd288a 100644
--- a/tag/map_codec_test.go
+++ b/tag/map_codec_test.go
@@ -22,45 +22,41 @@ import (
 	"testing"
 )
 
-var keys []Key
-
-func Test_EncodeDecode_Set(t *testing.T) {
+func TestEncodeDecode(t *testing.T) {
 	k1, _ := NewKey("k1")
 	k2, _ := NewKey("k2")
 	k3, _ := NewKey("k3 is very weird <>.,?/'\";:`~!@#$%^&*()_-+={[}]|\\")
 	k4, _ := NewKey("k4")
 
-	type pair struct {
+	type keyValue struct {
 		k Key
 		v string
 	}
 
-	type testCase struct {
+	testCases := []struct {
 		label string
-		pairs []pair
-	}
-
-	testCases := []testCase{
+		pairs []keyValue
+	}{
 		{
 			"0",
-			[]pair{},
+			[]keyValue{},
 		},
 		{
 			"1",
-			[]pair{
+			[]keyValue{
 				{k1, "v1"},
 			},
 		},
 		{
 			"2",
-			[]pair{
+			[]keyValue{
 				{k1, "v1"},
 				{k2, "v2"},
 			},
 		},
 		{
 			"3",
-			[]pair{
+			[]keyValue{
 				{k1, "v1"},
 				{k2, "v2"},
 				{k3, "v3"},
@@ -68,7 +64,7 @@ func Test_EncodeDecode_Set(t *testing.T) {
 		},
 		{
 			"4",
-			[]pair{
+			[]keyValue{
 				{k1, "v1"},
 				{k2, "v2"},
 				{k3, "v3"},
@@ -86,16 +82,16 @@ func Test_EncodeDecode_Set(t *testing.T) {
 		if err != nil {
 			t.Errorf("%v: NewMap = %v", tc.label, err)
 		}
+
 		encoded := Encode(ts)
 		decoded, err := Decode(encoded)
-
 		if err != nil {
 			t.Errorf("%v: decoding encoded tag map failed: %v", tc.label, err)
 		}
 
-		got := make([]pair, 0)
+		got := make([]keyValue, 0)
 		for k, v := range decoded.m {
-			got = append(got, pair{k, string(v)})
+			got = append(got, keyValue{k, string(v)})
 		}
 		want := tc.pairs
 
@@ -107,3 +103,52 @@ func Test_EncodeDecode_Set(t *testing.T) {
 		}
 	}
 }
+
+func TestDecode(t *testing.T) {
+	k1, _ := NewKey("k1")
+	m, _ := NewMap(context.Background(), Insert(k1, "v1"))
+
+	tests := []struct {
+		name    string
+		bytes   []byte
+		want    *Map
+		wantErr bool
+	}{
+		{
+			name:    "valid",
+			bytes:   []byte{0, 0, 2, 107, 49, 2, 118, 49},
+			want:    m,
+			wantErr: false,
+		},
+		{
+			name:    "non-ascii key",
+			bytes:   []byte{0, 0, 2, 107, 49, 2, 118, 49, 0, 2, 107, 25, 2, 118, 49},
+			want:    nil,
+			wantErr: true,
+		},
+		{
+			name:    "non-ascii value",
+			bytes:   []byte{0, 0, 2, 107, 49, 2, 118, 49, 0, 2, 107, 50, 2, 118, 25},
+			want:    nil,
+			wantErr: true,
+		},
+		{
+			name:    "long value",
+			bytes:   []byte{0, 0, 2, 107, 49, 2, 118, 49, 0, 2, 107, 50, 172, 2, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97},
+			want:    nil,
+			wantErr: true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := Decode(tt.bytes)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("Decode() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("Decode() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
diff --git a/trace/sampling_test.go b/trace/sampling_test.go
new file mode 100644
index 0000000..5953a8f
--- /dev/null
+++ b/trace/sampling_test.go
@@ -0,0 +1,46 @@
+// Copyright 2017, OpenCensus Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package trace
+
+import "testing"
+import "reflect"
+
+func TestSetDefaultSampler(t *testing.T) {
+	tests := []struct {
+		name    string
+		sampler Sampler
+		want    Sampler
+	}{
+		{
+			name:    "when the sampler is set to nil, the default sampler should be used",
+			sampler: nil,
+			want:    ProbabilitySampler(defaultSamplingProbability),
+		},
+		{
+			name:    "setting a NeverSample updates the sampler",
+			sampler: NeverSample(),
+			want:    NeverSample(),
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			SetDefaultSampler(tt.sampler)
+			if !reflect.DeepEqual(defaultSampler, tt.want) {
+				t.Errorf("%q. SetDefaultSampler() = %v, want %v", tt.name, defaultSampler, tt.want)
+			}
+			SetDefaultSampler(nil) // Need to reset the sampler between each test
+		})
+	}
+}