1 files changed, 18 insertions, 13 deletions
diff --git a/src/cmd/trace/trace_test.go b/src/cmd/trace/trace_test.go
index aaffda87f0..aff3863802 100644
--- a/src/cmd/trace/trace_test.go
+++ b/src/cmd/trace/trace_test.go
@@ -6,6 +6,7 @@ package main
 
 import (
 	"internal/trace"
+	"io/ioutil"
 	"strings"
 	"testing"
 )
@@ -68,12 +69,10 @@ func TestGoroutineCount(t *testing.T) {
 		endTime: int64(1<<63 - 1),
 	}
 
-	// If the counts drop below 0, generateTrace will return an error.
-	viewerData, err := generateTrace(params)
-	if err != nil {
-		t.Fatalf("generateTrace failed: %v", err)
-	}
-	for _, ev := range viewerData.Events {
+	// Use the default viewerDataTraceConsumer but replace
+	// consumeViewerEvent to intercept the ViewerEvents for testing.
+	c := viewerDataTraceConsumer(ioutil.Discard, 0, 1<<63-1)
+	c.consumeViewerEvent = func(ev *ViewerEvent, _ bool) {
 		if ev.Name == "Goroutines" {
 			cnt := ev.Arg.(*goroutineCountersArg)
 			if cnt.Runnable+cnt.Running > 2 {
@@ -82,6 +81,11 @@ func TestGoroutineCount(t *testing.T) {
 			t.Logf("read %+v %+v", ev, cnt)
 		}
 	}
+
+	// If the counts drop below 0, generateTrace will return an error.
+	if err := generateTrace(params, c); err != nil {
+		t.Fatalf("generateTrace failed: %v", err)
+	}
 }
 
 func TestGoroutineFilter(t *testing.T) {
@@ -120,8 +124,8 @@ func TestGoroutineFilter(t *testing.T) {
 		gs:      map[uint64]bool{10: true},
 	}
 
-	_, err = generateTrace(params)
-	if err != nil {
+	c := viewerDataTraceConsumer(ioutil.Discard, 0, 1<<63-1)
+	if err := generateTrace(params, c); err != nil {
 		t.Fatalf("generateTrace failed: %v", err)
 	}
 }
@@ -152,17 +156,18 @@ func TestPreemptedMarkAssist(t *testing.T) {
 		endTime: int64(1<<63 - 1),
 	}
 
-	viewerData, err := generateTrace(params)
-	if err != nil {
-		t.Fatalf("generateTrace failed: %v", err)
-	}
+	c := viewerDataTraceConsumer(ioutil.Discard, 0, 1<<63-1)
 
 	marks := 0
-	for _, ev := range viewerData.Events {
+	c.consumeViewerEvent = func(ev *ViewerEvent, _ bool) {
 		if strings.Contains(ev.Name, "MARK ASSIST") {
 			marks++
 		}
 	}
+	if err := generateTrace(params, c); err != nil {
+		t.Fatalf("generateTrace failed: %v", err)
+	}
+
 	if marks != 2 {
 		t.Errorf("Got %v MARK ASSIST events, want %v", marks, 2)
 	}