From 9c1758b5931f9d55ee3080d943018aee3adbb0e0 Mon Sep 17 00:00:00 2001
From: bergquist <carl.bergquist@gmail.com>
Date: Wed, 13 Jun 2018 09:18:53 +0200
Subject: [PATCH] bus: Dispatch now passes empty ctx if handler require it

---
 pkg/bus/bus.go      | 16 +++++++++++++---
 pkg/bus/bus_test.go | 17 +++++++++--------
 2 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/pkg/bus/bus.go b/pkg/bus/bus.go
index 1259211cddc..9cd3f116172 100644
--- a/pkg/bus/bus.go
+++ b/pkg/bus/bus.go
@@ -96,13 +96,23 @@ func (b *InProcBus) DispatchCtx(ctx context.Context, msg Msg) error {
 func (b *InProcBus) Dispatch(msg Msg) error {
 	var msgName = reflect.TypeOf(msg).Elem().Name()
 
-	var handler = b.handlers[msgName]
+	var handler = b.handlersWithCtx[msgName]
+	withCtx := true
+
+	if handler == nil {
+		withCtx = false
+		handler = b.handlers[msgName]
+	}
+
 	if handler == nil {
 		return ErrHandlerNotFound
 	}
 
-	var params = make([]reflect.Value, 1)
-	params[0] = reflect.ValueOf(msg)
+	var params = []reflect.Value{}
+	if withCtx {
+		params = append(params, reflect.ValueOf(context.Background()))
+	}
+	params = append(params, reflect.ValueOf(msg))
 
 	ret := reflect.ValueOf(handler).Call(params)
 	err := ret[0].Interface()
diff --git a/pkg/bus/bus_test.go b/pkg/bus/bus_test.go
index 83a0d7190ea..9f41a5154df 100644
--- a/pkg/bus/bus_test.go
+++ b/pkg/bus/bus_test.go
@@ -34,7 +34,6 @@ func TestDispatchCtxCanUseNormalHandlers(t *testing.T) {
 	}
 
 	bus.AddHandler(handler)
-	bus.AddHandlerCtx(handlerWithCtx)
 
 	t.Run("when a normal handler is registered", func(t *testing.T) {
 		bus.Dispatch(&testQuery{})
@@ -42,15 +41,17 @@ func TestDispatchCtxCanUseNormalHandlers(t *testing.T) {
 		if handlerCallCount != 1 {
 			t.Errorf("Expected normal handler to be called 1 time. was called %d", handlerCallCount)
 		}
+
+		t.Run("when a ctx handler is registered", func(t *testing.T) {
+			bus.AddHandlerCtx(handlerWithCtx)
+			bus.Dispatch(&testQuery{})
+
+			if handlerWithCtxCallCount != 1 {
+				t.Errorf("Expected ctx handler to be called 1 time. was called %d", handlerWithCtxCallCount)
+			}
+		})
 	})
 
-	t.Run("when a ctx handler is registered", func(t *testing.T) {
-		bus.DispatchCtx(context.Background(), &testQuery{})
-
-		if handlerWithCtxCallCount != 1 {
-			t.Errorf("Expected ctx handler to be called 1 time. was called %d", handlerWithCtxCallCount)
-		}
-	})
 }
 
 func TestQueryHandlerReturnsError(t *testing.T) {