Refactor router dispatcher tests to include context parameter

ksysoev · Jun 9, 2024 · 154f2f7 · 154f2f7
1 parent df2bee9
commit 154f2f7
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/dispatch/router_dipatcher_test.go b/dispatch/router_dipatcher_test.go
@@ -69,6 +69,8 @@ func TestRouterDispatcher_DispatchDefaultBackend(t *testing.T) {
 	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
@@ -89,6 +91,8 @@ func TestRouterDispatcher_DispatchByRoutingKey(t *testing.T) {
 	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
@@ -110,6 +114,8 @@ func TestRouterDispatcher_DispatchWrongRequest(t *testing.T) {
 	dispatcher := NewRouterDispatcher(defaultBackend, parser)
 
 	conn := mocks.NewMockConnection(t)
+	conn.EXPECT().Context().Return(context.Background())
+
 	data := []byte("test data")
 
 	dispatcher.Dispatch(conn, wasabi.MsgTypeText, data)
@@ -124,6 +130,8 @@ func TestRouterDispatcher_DispatchErrorHandlingRequest(t *testing.T) {
 	dispatcher := NewRouterDispatcher(defaultBackend, parser)
 
 	conn := mocks.NewMockConnection(t)
+	conn.EXPECT().Context().Return(context.Background())
+
 	data := []byte("test data")
 
 	routingKey := "key1"