From cb0f217e94d2c5bfb03e76ef6a8842cbec8b1a75 Mon Sep 17 00:00:00 2001
From: Kirill Sysoev <k.sysoev@me.com>
Date: Sat, 18 May 2024 15:42:54 +0800
Subject: [PATCH] Adds request factopry usage to test cases and code examples

---
 backend/ws_test.go           | 52 +++++++++++++++++++++++++++++++-----
 examples/passthrough/main.go | 15 ++++++++++-
 2 files changed, 59 insertions(+), 8 deletions(-)

diff --git a/backend/ws_test.go b/backend/ws_test.go
index 9b45c6f..1a6d083 100644
--- a/backend/ws_test.go
+++ b/backend/ws_test.go
@@ -8,6 +8,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/ksysoev/wasabi"
 	"github.com/ksysoev/wasabi/mocks"
 	"nhooyr.io/websocket"
 )
@@ -47,7 +48,9 @@ var wsHandlerEcho = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request
 
 func TestNewWSBackend(t *testing.T) {
 	url := "ws://example.com"
-	b := NewWSBackend(url)
+	b := NewWSBackend(url, func(r wasabi.Request) (websocket.MessageType, []byte, error) {
+		return websocket.MessageText, []byte("Hello, world!"), nil
+	})
 
 	if b.URL != url {
 		t.Errorf("Expected URL to be %q, but got %q", url, b.URL)
@@ -67,7 +70,9 @@ func TestGetConnectionExistingConnection(t *testing.T) {
 	defer server.Close()
 	url := "ws://" + server.Listener.Addr().String()
 
-	b := NewWSBackend(url)
+	b := NewWSBackend(url, func(r wasabi.Request) (websocket.MessageType, []byte, error) {
+		return websocket.MessageText, []byte("Hello, world!"), nil
+	})
 
 	conn := mocks.NewMockConnection(t)
 	conn.EXPECT().ID().Return("connection1")
@@ -90,7 +95,9 @@ func TestGetConnectionNewConnection(t *testing.T) {
 	defer server.Close()
 	url := "ws://" + server.Listener.Addr().String()
 
-	b := NewWSBackend(url)
+	b := NewWSBackend(url, func(r wasabi.Request) (websocket.MessageType, []byte, error) {
+		return websocket.MessageText, []byte("Hello, world!"), nil
+	})
 
 	conn := mocks.NewMockConnection(t)
 	conn.EXPECT().ID().Return("connection1")
@@ -115,7 +122,9 @@ func TestGetConnectionDialError(t *testing.T) {
 	url := "ws://" + server.Listener.Addr().String()
 	server.Close()
 
-	b := NewWSBackend(url)
+	b := NewWSBackend(url, func(r wasabi.Request) (websocket.MessageType, []byte, error) {
+		return websocket.MessageText, []byte("Hello, world!"), nil
+	})
 
 	conn := mocks.NewMockConnection(t)
 	conn.EXPECT().ID().Return("connection1")
@@ -151,7 +160,9 @@ func TestWSBackend_Handle(t *testing.T) {
 	r.EXPECT().Data().Return([]byte("Hello, world!"))
 	r.EXPECT().Context().Return(context.Background())
 
-	b := NewWSBackend(url)
+	b := NewWSBackend(url, func(r wasabi.Request) (websocket.MessageType, []byte, error) {
+		return websocket.MessageText, r.Data(), nil
+	})
 
 	err := b.Handle(conn, r)
 
@@ -173,7 +184,9 @@ func TestWSBackend_Handle_FailToConnect(t *testing.T) {
 
 	r := mocks.NewMockRequest(t)
 
-	b := NewWSBackend(url)
+	b := NewWSBackend(url, func(r wasabi.Request) (websocket.MessageType, []byte, error) {
+		return websocket.MessageText, []byte("Hello, world!"), nil
+	})
 
 	err := b.Handle(conn, r)
 
@@ -196,7 +209,9 @@ func TestWSBackend_Handle_CloseConnection(t *testing.T) {
 
 	conn.EXPECT().Close(websocket.StatusNormalClosure, "").Return(nil)
 
-	b := NewWSBackend(url)
+	b := NewWSBackend(url, func(r wasabi.Request) (websocket.MessageType, []byte, error) {
+		return websocket.MessageText, []byte("Hello, world!"), nil
+	})
 
 	wsConn, resp, err := websocket.Dial(ctx, url, nil)
 	if err != nil {
@@ -222,3 +237,26 @@ func TestWSBackend_Handle_CloseConnection(t *testing.T) {
 		t.Error("Expected connection to be closed")
 	}
 }
+
+func TestWSBackend_RequestFactory_Error(t *testing.T) {
+	server := httptest.NewServer(wsHandlerEcho)
+	url := "ws://" + server.Listener.Addr().String()
+
+	defer server.Close()
+
+	conn := mocks.NewMockConnection(t)
+	conn.EXPECT().ID().Return("connection1")
+	conn.EXPECT().Context().Return(context.Background())
+
+	r := mocks.NewMockRequest(t)
+
+	b := NewWSBackend(url, func(r wasabi.Request) (websocket.MessageType, []byte, error) {
+		return 0, nil, io.EOF
+	})
+
+	err := b.Handle(conn, r)
+
+	if err == nil {
+		t.Errorf("Expected error, but got nil")
+	}
+}
diff --git a/examples/passthrough/main.go b/examples/passthrough/main.go
index d047be5..e5ea65c 100644
--- a/examples/passthrough/main.go
+++ b/examples/passthrough/main.go
@@ -22,7 +22,20 @@ func main() {
 
 	slog.LogAttrs(context.Background(), slog.LevelDebug, "")
 
-	backend := backend.NewWSBackend("wss://ws.derivws.com/websockets/v3?app_id=1089")
+	backend := backend.NewWSBackend(
+		"wss://ws.derivws.com/websockets/v3?app_id=1089",
+		func(r wasabi.Request) (wasabi.MessageType, []byte, error) {
+			switch r.RoutingKey() {
+			case "text":
+				return wasabi.MsgTypeText, r.Data(), nil
+			case "binary":
+				return wasabi.MsgTypeBinary, r.Data(), nil
+			default:
+				var t wasabi.MessageType
+				return t, nil, fmt.Errorf("unsupported request type: %s", r.RoutingKey())
+			}
+		},
+	)
 
 	dispatcher := dispatch.NewRouterDispatcher(backend, func(conn wasabi.Connection, msgType wasabi.MessageType, data []byte) wasabi.Request {
 		return dispatch.NewRawRequest(conn.Context(), msgType, data)