Merge pull request #79 from ksysoev/reqest_context_to_dispatcher

"Refactor request parser and router dispatcher to include context parameter"
ksysoev · Jun 9, 2024 · ba32846 · ba32846
2 parents 43e7c02 + 0049d3f
commit ba32846
Show file tree

Hide file tree

Showing 7 changed files with 39 additions and 18 deletions.
diff --git a/dispatch/common.go b/dispatch/common.go
@@ -1,6 +1,10 @@
 package dispatch
 
-import "github.com/ksysoev/wasabi"
+import (
+	"context"
+
+	"github.com/ksysoev/wasabi"
+)
 
 // RequestMiddlewere is interface for request middleweres
 type RequestMiddlewere func(next wasabi.RequestHandler) wasabi.RequestHandler
@@ -13,4 +17,4 @@ func (f RequestHandlerFunc) Handle(conn wasabi.Connection, req wasabi.Request) e
 	return f(conn, req)
 }
 
-type RequestParser func(conn wasabi.Connection, msgType wasabi.MessageType, data []byte) wasabi.Request
+type RequestParser func(conn wasabi.Connection, ctx context.Context, msgType wasabi.MessageType, data []byte) wasabi.Request
diff --git a/dispatch/router_dipatcher.go b/dispatch/router_dipatcher.go
@@ -44,7 +44,7 @@ func (d *RouterDispatcher) AddBackend(backend wasabi.RequestHandler, routingKeys
 // determining the appropriate backend, and handling the request using middleware.
 // If an error occurs during handling, it is logged.
 func (d *RouterDispatcher) Dispatch(conn wasabi.Connection, msgType wasabi.MessageType, data []byte) {
-	req := d.parser(conn, msgType, data)
+	req := d.parser(conn, conn.Context(), msgType, data)
 
 	if req == nil {
 		return

diff --git a/dispatch/router_dipatcher_test.go b/dispatch/router_dipatcher_test.go
@@ -1,6 +1,7 @@
 package dispatch
 
 import (
+	"context"
 	"fmt"
 	"testing"
 
@@ -11,7 +12,7 @@ import (
 func TestNewRouterDispatcher(t *testing.T) {
 	defaultBackend := mocks.NewMockBackend(t)
 
-	parser := func(_ wasabi.Connection, _ wasabi.MessageType, _ []byte) wasabi.Request {
+	parser := func(_ wasabi.Connection, _ context.Context, _ wasabi.MessageType, _ []byte) wasabi.Request {
 		return mocks.NewMockRequest(t)
 	}
 
@@ -27,7 +28,7 @@ func TestNewRouterDispatcher(t *testing.T) {
 }
 func TestRouterDispatcher_AddBackend(t *testing.T) {
 	defaultBackend := mocks.NewMockBackend(t)
-	parser := func(_ wasabi.Connection, _ wasabi.MessageType, _ []byte) wasabi.Request {
+	parser := func(_ wasabi.Connection, _ context.Context, _ wasabi.MessageType, _ []byte) wasabi.Request {
 		return mocks.NewMockRequest(t)
 	}
 	dispatcher := NewRouterDispatcher(defaultBackend, parser)
@@ -62,10 +63,14 @@ func TestRouterDispatcher_DispatchDefaultBackend(t *testing.T) {
 	defaultBackend := mocks.NewMockBackend(t)
 
 	req := mocks.NewMockRequest(t)
-	parser := func(_ wasabi.Connection, _ wasabi.MessageType, _ []byte) wasabi.Request { return req }
+	parser := func(_ wasabi.Connection, _ context.Context, _ wasabi.MessageType, _ []byte) wasabi.Request {
+		return req
+	}
 	dispatcher := NewRouterDispatcher(defaultBackend, parser)
 
 	conn := mocks.NewMockConnection(t)
+	conn.EXPECT().Context().Return(context.Background())
+
 	data := []byte("test data")
 
 	// Test case 1: Request with existing routing key
@@ -80,10 +85,14 @@ func TestRouterDispatcher_DispatchDefaultBackend(t *testing.T) {
 func TestRouterDispatcher_DispatchByRoutingKey(t *testing.T) {
 	defaultBackend := mocks.NewMockBackend(t)
 	req := mocks.NewMockRequest(t)
-	parser := func(_ wasabi.Connection, _ wasabi.MessageType, _ []byte) wasabi.Request { return req }
+	parser := func(_ wasabi.Connection, _ context.Context, _ wasabi.MessageType, _ []byte) wasabi.Request {
+		return req
+	}
 	dispatcher := NewRouterDispatcher(defaultBackend, parser)
 
 	conn := mocks.NewMockConnection(t)
+	conn.EXPECT().Context().Return(context.Background())
+
 	data := []byte("test data")
 
 	// Test case 1: Request with existing routing key
@@ -99,10 +108,14 @@ func TestRouterDispatcher_DispatchByRoutingKey(t *testing.T) {
 
 func TestRouterDispatcher_DispatchWrongRequest(t *testing.T) {
 	defaultBackend := mocks.NewMockBackend(t)
-	parser := func(_ wasabi.Connection, _ wasabi.MessageType, _ []byte) wasabi.Request { return nil }
+	parser := func(_ wasabi.Connection, _ context.Context, _ wasabi.MessageType, _ []byte) wasabi.Request {
+		return nil
+	}
 	dispatcher := NewRouterDispatcher(defaultBackend, parser)
 
 	conn := mocks.NewMockConnection(t)
+	conn.EXPECT().Context().Return(context.Background())
+
 	data := []byte("test data")
 
 	dispatcher.Dispatch(conn, wasabi.MsgTypeText, data)
@@ -111,10 +124,14 @@ func TestRouterDispatcher_DispatchWrongRequest(t *testing.T) {
 func TestRouterDispatcher_DispatchErrorHandlingRequest(t *testing.T) {
 	defaultBackend := mocks.NewMockBackend(t)
 	req := mocks.NewMockRequest(t)
-	parser := func(_ wasabi.Connection, _ wasabi.MessageType, _ []byte) wasabi.Request { return req }
+	parser := func(_ wasabi.Connection, _ context.Context, _ wasabi.MessageType, _ []byte) wasabi.Request {
+		return req
+	}
 	dispatcher := NewRouterDispatcher(defaultBackend, parser)
 
 	conn := mocks.NewMockConnection(t)
+	conn.EXPECT().Context().Return(context.Background())
+
 	data := []byte("test data")
 
 	routingKey := "key1"
@@ -128,7 +145,7 @@ func TestRouterDispatcher_DispatchErrorHandlingRequest(t *testing.T) {
 }
 func TestRouterDispatcher_Use(t *testing.T) {
 	defaultBackend := mocks.NewMockBackend(t)
-	parser := func(_ wasabi.Connection, _ wasabi.MessageType, _ []byte) wasabi.Request {
+	parser := func(_ wasabi.Connection, _ context.Context, _ wasabi.MessageType, _ []byte) wasabi.Request {
 		return mocks.NewMockRequest(t)
 	}
 	dispatcher := NewRouterDispatcher(defaultBackend, parser)
@@ -149,7 +166,7 @@ func TestRouterDispatcher_UseMiddleware(t *testing.T) {
 	defaultBackend := mocks.NewMockBackend(t)
 	defaultBackend.EXPECT().Handle(mockConn, mockReq).Return(testError)
 
-	parser := func(_ wasabi.Connection, _ wasabi.MessageType, _ []byte) wasabi.Request {
+	parser := func(_ wasabi.Connection, _ context.Context, _ wasabi.MessageType, _ []byte) wasabi.Request {
 		return mocks.NewMockRequest(t)
 	}
 	dispatcher := NewRouterDispatcher(defaultBackend, parser)

diff --git a/examples/echo/main.go b/examples/echo/main.go
@@ -25,8 +25,8 @@ func main() {
 		return conn.Send(wasabi.MsgTypeText, req.Data())
 	})
 
-	dispatcher := dispatch.NewRouterDispatcher(backend, func(conn wasabi.Connection, msgType wasabi.MessageType, data []byte) wasabi.Request {
-		return dispatch.NewRawRequest(conn.Context(), msgType, data)
+	dispatcher := dispatch.NewRouterDispatcher(backend, func(conn wasabi.Connection, ctx context.Context, msgType wasabi.MessageType, data []byte) wasabi.Request {
+		return dispatch.NewRawRequest(ctx, msgType, data)
 	})
 	channel := channel.NewChannel("/", dispatcher, channel.NewConnectionRegistry(), channel.WithOriginPatterns("*"))
 

diff --git a/examples/http_backend/main.go b/examples/http_backend/main.go
@@ -30,8 +30,8 @@ func main() {
 		return httpReq, nil
 	})
 
-	dispatcher := dispatch.NewRouterDispatcher(backend, func(conn wasabi.Connection, msgType wasabi.MessageType, data []byte) wasabi.Request {
-		return dispatch.NewRawRequest(conn.Context(), msgType, data)
+	dispatcher := dispatch.NewRouterDispatcher(backend, func(conn wasabi.Connection, ctx context.Context, msgType wasabi.MessageType, data []byte) wasabi.Request {
+		return dispatch.NewRawRequest(ctx, msgType, data)
 	})
 
 	channel := channel.NewChannel("/", dispatcher, channel.NewConnectionRegistry(), channel.WithOriginPatterns("*"))

diff --git a/examples/passthrough/main.go b/examples/passthrough/main.go
@@ -37,7 +37,7 @@ func main() {
 		},
 	)
 
-	dispatcher := dispatch.NewRouterDispatcher(backend, func(conn wasabi.Connection, msgType wasabi.MessageType, data []byte) wasabi.Request {
+	dispatcher := dispatch.NewRouterDispatcher(backend, func(conn wasabi.Connection, ctx context.Context, msgType wasabi.MessageType, data []byte) wasabi.Request {
 		return dispatch.NewRawRequest(conn.Context(), msgType, data)
 	})
 	channel := channel.NewChannel("/", dispatcher, channel.NewConnectionRegistry(), channel.WithOriginPatterns("*"))

diff --git a/tests/echo_test.go b/tests/echo_test.go
@@ -22,8 +22,8 @@ func TestEcho(t *testing.T) {
 		return conn.Send(wasabi.MsgTypeText, req.Data())
 	})
 
-	dispatcher := dispatch.NewRouterDispatcher(backend, func(conn wasabi.Connection, msgType wasabi.MessageType, data []byte) wasabi.Request {
-		return dispatch.NewRawRequest(conn.Context(), msgType, data)
+	dispatcher := dispatch.NewRouterDispatcher(backend, func(conn wasabi.Connection, ctx context.Context, msgType wasabi.MessageType, data []byte) wasabi.Request {
+		return dispatch.NewRawRequest(ctx, msgType, data)
 	})
 	ch := channel.NewChannel("/", dispatcher, channel.NewConnectionRegistry(), channel.WithOriginPatterns("*"))