playground/socket: deflake test, remove sleep

Fixes golang/go#18275

Change-Id: Iedb0a21a3fe0b383a2e51fc92076d7fd005d71e6
Reviewed-on: https://go-review.googlesource.com/34277
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Ian Lance Taylor <iant@golang.org>
diff --git a/playground/socket/socket.go b/playground/socket/socket.go
index 527c075..0c275bf 100644
--- a/playground/socket/socket.go
+++ b/playground/socket/socket.go
@@ -175,7 +175,7 @@
 	)
 	go func() {
 		defer close(done)
-		for m := range buffer(limiter(out, p)) {
+		for m := range buffer(limiter(out, p), time.After) {
 			m.Id = id
 			dest <- m
 		}
@@ -259,7 +259,8 @@
 // coalesced; when a message of a different kind is received, any buffered
 // messages are flushed. When the given channel is closed, buffer flushes the
 // remaining buffered messages and closes the returned channel.
-func buffer(in <-chan *Message) <-chan *Message {
+// The timeAfter func should be time.After. It exists for testing.
+func buffer(in <-chan *Message, timeAfter func(time.Duration) <-chan time.Time) <-chan *Message {
 	out := make(chan *Message)
 	go func() {
 		defer close(out)
@@ -292,7 +293,7 @@
 					flush()
 					kind = m.Kind
 					if tc == nil {
-						tc = time.After(msgDelay)
+						tc = timeAfter(msgDelay)
 					}
 				}
 				buf = append(buf, m.Body...)
diff --git a/playground/socket/socket_test.go b/playground/socket/socket_test.go
index 5dd2815..0ceb250 100644
--- a/playground/socket/socket_test.go
+++ b/playground/socket/socket_test.go
@@ -10,24 +10,28 @@
 )
 
 func TestBuffer(t *testing.T) {
+	afterChan := make(chan time.Time)
 	ch := make(chan *Message)
 	go func() {
 		ch <- &Message{Kind: "err", Body: "a"}
 		ch <- &Message{Kind: "err", Body: "b"}
 		ch <- &Message{Kind: "out", Body: "1"}
 		ch <- &Message{Kind: "out", Body: "2"}
-		time.Sleep(msgDelay * 2)
+		afterChan <- time.Time{} // value itself doesn't matter
 		ch <- &Message{Kind: "out", Body: "3"}
 		ch <- &Message{Kind: "out", Body: "4"}
 		close(ch)
 	}()
 
 	var ms []*Message
-	for m := range buffer(ch) {
+	timeAfter := func(d time.Duration) <-chan time.Time {
+		return afterChan
+	}
+	for m := range buffer(ch, timeAfter) {
 		ms = append(ms, m)
 	}
 	if len(ms) != 3 {
-		t.Fatalf("got %v messages, want 2", len(ms))
+		t.Fatalf("got %v messages, want 3", len(ms))
 	}
 	if g, w := ms[0].Body, "ab"; g != w {
 		t.Errorf("message 0 body = %q, want %q", g, w)