diff --git a/dispatch/common.go b/dispatch/common.go index b4b8d30..665d3d7 100644 --- 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 index d8ac17b..fe50ad4 100644 --- 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 index 97c7065..05b8421 100644 --- 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 index 6db6aa5..7d7b74e 100644 --- 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 index d71bd14..4325076 100644 --- 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 index e5ea65c..12f04e8 100644 --- 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 index 3cb881d..0234044 100644 --- 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("*"))