Handler: Changed ChanFor API, fixed race conditions in tests

handler/handler.go

@@ -75,9 +75,10 @@ func (h *Handler) Call(ev interface{}) {
 	}
 }
 
-func (h *Handler) WaitFor(
-	ctx context.Context, fn func(interface{}) bool) interface{} {
-
+// WaitFor blocks until there's an event. It's advised to use ChanFor instead,
+// as WaitFor may skip some events if it's not ran fast enough after the event
+// arrived.
+func (h *Handler) WaitFor(ctx context.Context, fn func(interface{}) bool) interface{} {
 	var result = make(chan interface{})
 
 	cancel := h.AddHandler(func(v interface{}) {
@@ -95,22 +96,29 @@ func (h *Handler) WaitFor(
 	}
 }
 
-func (h *Handler) ChanFor(fn func(interface{}) bool) <-chan interface{} {
-	var result = make(chan interface{})
+// ChanFor returns a channel that would receive all incoming events that match
+// the callback given. The cancel() function removes the handler and drops all
+// hanging goroutines.
+func (h *Handler) ChanFor(fn func(interface{}) bool) (out <-chan interface{}, cancel func()) {
+	result := make(chan interface{})
+	closer := make(chan struct{})
 
-	cancel := h.AddHandler(func(v interface{}) {
+	removeHandler := h.AddHandler(func(v interface{}) {
 		if fn(v) {
-			result <- v
+			select {
+			case result <- v:
+			case <-closer:
+			}
 		}
 	})
 
-	var recv = make(chan interface{})
-	go func() {
-		recv <- <-result
-		cancel()
-	}()
+	cancel = func() {
+		removeHandler()
+		close(closer)
+	}
+	out = result
 
-	return recv
+	return
 }
 
 func (h *Handler) AddHandler(handler interface{}) (rm func()) {

handler/handler_test.go

@@ -125,7 +125,7 @@ func TestHandlerInterface(t *testing.T) {
 }
 
 func TestHandlerWait(t *testing.T) {
-	inc := make(chan interface{})
+	inc := make(chan interface{}, 1)
 
 	h := New()
 
@@ -151,19 +151,14 @@ func TestHandlerWait(t *testing.T) {
 		})
 	}()
 
-	var recv interface{}
-	var done = make(chan struct{})
-	go func() {
-		recv = <-inc
-		done <- struct{}{}
-	}()
+	// Wait for WaitFor to add its handler:
+	time.Sleep(time.Millisecond)
 
 	for _, ev := range evs {
-		time.Sleep(1)
 		h.Call(ev)
 	}
 
-	<-done
+	recv := <-inc
 	if recv != wanted {
 		t.Fatal("Unexpected receive:", recv)
 	}
@@ -195,14 +190,15 @@ func TestHandlerChan(t *testing.T) {
 		wanted,
 	}
 
-	inc := h.ChanFor(func(v interface{}) bool {
+	inc, cancel := h.ChanFor(func(v interface{}) bool {
 		tp, ok := v.(*gateway.TypingStartEvent)
 		if !ok {
 			return false
 		}
 
 		return tp.ChannelID == wanted.ChannelID
 	})
+	defer cancel()
 
 	for _, ev := range evs {
 		h.Call(ev)