Merge pull request #35 from danschultzer/allow-to-for-websocket-init

Allow `:to` option for `TestServer.websocket_init/3`

Author: danschultzer

CHANGELOG.md

@@ -1,5 +1,9 @@
 # Changelog
 
+## v0.1.19 (TBA)
+
+- Allow `:to` plug to be set for `TestServer.websocket_init/3` for handshake
+
 ## v0.1.18 (2024-12-29)
 
 - Limit number of Bandit acceptors to 1 to improve performance

lib/test_server.ex

@@ -353,17 +353,21 @@ defmodule TestServer do
   """
   @spec add(pid(), binary(), keyword()) :: :ok
   def add(instance, uri, options) when is_pid(instance) and is_binary(uri) and is_list(options) do
-    instance_alive!(instance)
-
-    [_first_module_entry | stacktrace] = get_stacktrace()
-
     options = Keyword.put_new(options, :to, &default_response_handler/1)
 
-    {:ok, _route} = Instance.register(instance, {:plug_router_to, {uri, options, stacktrace}})
+    {:ok, _route} = register_route(instance, uri, options)
 
     :ok
   end
 
+  defp register_route(instance, uri, options) do
+    instance_alive!(instance)
+
+    [_register_route, _first_module_entry | stacktrace] = get_stacktrace()
+
+    Instance.register(instance, {:plug_router_to, {uri, options, stacktrace}})
+  end
+
   defp get_stacktrace do
     {:current_stacktrace, [{Process, :info, _, _} | stacktrace]} =
       Process.info(self(), :current_stacktrace)
@@ -480,6 +484,10 @@ defmodule TestServer do
   @doc """
   Adds a websocket route to current test server.
 
+  The `:to` option can be overridden the same way as for `add/2`, and will be
+  called during the HTTP handshake. If the `conn.state` is `:unset` the
+  websocket will be initiated otherwise response is returned as-is.
+
   ## Options
 
   Takes the same options as `add/2`, except `:to`.
@@ -499,6 +507,17 @@ defmodule TestServer do
       end)
 
       assert {:ok, _client} = WebSocketClient.start_link(TestServer.url("/ws?token=secret"))
+
+  `:to` option is also called during the HTTP handshake:
+
+      TestServer.websocket_init("/ws",
+        to: fn conn ->
+          Plug.Conn.send_resp(conn, 403, "Forbidden")
+        end
+      )
+
+      assert {:error, %WebSockex.RequestError{code: 403}} =
+              WebSocketClient.start_link(TestServer.url("/ws"))
   """
   @spec websocket_init(binary(), keyword()) :: {:ok, websocket_socket()}
   def websocket_init(uri, options) when is_binary(uri) do
@@ -519,16 +538,12 @@ defmodule TestServer do
   """
   @spec websocket_init(pid(), binary(), keyword()) :: {:ok, websocket_socket()}
   def websocket_init(instance, uri, options) do
-    instance_alive!(instance)
-
-    if Keyword.has_key?(options, :to), do: raise(ArgumentError, "`:to` is an invalid option")
-
-    [_first_module_entry | stacktrace] = get_stacktrace()
-
-    options = Keyword.put(options, :to, :websocket)
+    options =
+      options
+      |> Keyword.put(:websocket, true)
+      |> Keyword.put_new(:to, & &1)
 
-    {:ok, %{ref: ref}} =
-      Instance.register(instance, {:plug_router_to, {uri, options, stacktrace}})
+    {:ok, %{ref: ref}} = register_route(instance, uri, options)
 
     {:ok, {instance, ref}}
   end

lib/test_server/instance.ex

@@ -16,7 +16,7 @@ defmodule TestServer.Instance do
   @spec register(pid(), {:plug_router_to, {binary(), keyword(), TestServer.stacktrace()}}) ::
           {:ok, %{ref: reference()}}
   def register(instance, {:plug_router_to, {uri, options, stacktrace}}) do
-    options[:to] != :websocket && ensure_plug!(options[:to])
+    ensure_plug!(options[:to])
     options[:match] && ensure_function!(options[:match])
 
     GenServer.call(instance, {:register, {:plug_router_to, {uri, options, stacktrace}}})
@@ -319,16 +319,16 @@ defmodule TestServer.Instance do
       plugs -> plugs
     end
     |> Enum.reduce_while(conn, fn %{plug: plug, stacktrace: stacktrace}, conn ->
-      case try_run_plug(plug, stacktrace, conn) do
+      case try_run_plug(conn, plug, stacktrace) do
         {:ok, conn} -> {:cont, conn}
         {:error, error} -> {:halt, {:error, error}}
       end
     end)
   end
 
-  defp try_run_plug(plug, stacktrace, conn) do
-    plug
-    |> run_plug(conn)
+  defp try_run_plug(conn, plug, stacktrace) do
+    conn
+    |> run_plug(plug)
     |> check_halted!(plug, stacktrace)
   rescue
     error -> {:error, {error, __STACKTRACE__}}
@@ -345,11 +345,11 @@ defmodule TestServer.Instance do
 
   defp check_halted!(conn, _plug, _stacktrace), do: {:ok, conn}
 
-  defp run_plug(plug, conn) when is_function(plug) do
+  defp run_plug(conn, plug) when is_function(plug) do
     plug.(conn)
   end
 
-  defp run_plug(plug, conn) when is_atom(plug) do
+  defp run_plug(conn, plug) when is_atom(plug) do
     options = plug.init([])
     plug.call(conn, options)
   end
@@ -367,16 +367,12 @@ defmodule TestServer.Instance do
         {{:error, {:not_found, conn}}, state}
 
       index ->
+        %{to: plug, stacktrace: stacktrace} = route = Enum.at(state.routes, index)
+
         result =
-          case Enum.at(state.routes, index) do
-            %{to: :websocket, options: options} = route ->
-              websocket = {{self(), route.ref}, Keyword.get(options, :init_state)}
-              conn = Map.put(conn, :private, %{websocket: websocket})
-              {:ok, conn}
-
-            %{to: plug, stacktrace: stacktrace} ->
-              try_run_plug(plug, stacktrace, conn)
-          end
+          conn
+          |> maybe_put_websocket(route)
+          |> try_run_plug(plug, stacktrace)
 
         routes =
           List.update_at(state.routes, index, fn route ->
@@ -393,6 +389,17 @@ defmodule TestServer.Instance do
     error -> {:error, {error, __STACKTRACE__}}
   end
 
+  def maybe_put_websocket(conn, route) do
+    case route.options[:websocket] do
+      true ->
+        websocket = {{self(), route.ref}, Keyword.get(route.options, :init_state)}
+        Map.put(conn, :private, %{websocket: websocket})
+
+      _false ->
+        conn
+    end
+  end
+
   defp run_websocket_handlers({_instance, route_ref}, frame, websocket_state, state) do
     state.websocket_handlers
     |> Enum.map(&{&1.route_ref == route_ref, &1})

lib/test_server/plug.ex

@@ -8,7 +8,7 @@ defmodule TestServer.Plug do
 
   def call(conn, {http_server, args, instance}) do
     case Instance.dispatch(instance, {:plug, conn}) do
-      {:ok, %{private: %{websocket: {socket, state}}} = conn} ->
+      {:ok, %{state: :unset, private: %{websocket: {socket, state}}} = conn} ->
         :ok = Instance.put_websocket_connection(socket, http_server.get_socket_pid(conn))
         Plug.Conn.upgrade_adapter(conn, :websocket, {http_server, {socket, state}, args})
 

test/test_server_test.exs

@@ -558,8 +558,12 @@ defmodule TestServerTest do
       end
 
       test "invalid options" do
-        assert_raise ArgumentError, "`:to` is an invalid option", fn ->
-          TestServer.websocket_init("/", to: MyPlug)
+        assert_raise BadFunctionError, ~r/expected a function, got: :invalid/, fn ->
+          TestServer.websocket_init("/", to: :invalid)
+        end
+
+        assert_raise BadFunctionError, ~r/expected a function, got: :invalid/, fn ->
+          TestServer.websocket_init("/", match: :invalid)
         end
 
         assert_raise BadFunctionError, ~r/expected a function, got: :invalid/, fn ->
@@ -577,6 +581,31 @@ defmodule TestServerTest do
                        TestServer.websocket_init("/ws")
                      end
       end
+
+      test "with handshake callback function with set conn" do
+        assert {:ok, _socket} =
+                 TestServer.websocket_init("/ws",
+                   to: fn conn ->
+                     Plug.Conn.resp(conn, 403, "Forbidden")
+                   end
+                 )
+
+        assert {:error, %WebSockex.RequestError{code: 403}} =
+                 WebSocketClient.start_link(TestServer.url("/ws"))
+      end
+
+      test "with handshake callback function with unset conn" do
+        assert {:ok, _socket} =
+                 TestServer.websocket_init("/ws",
+                   to: fn conn ->
+                     assert Plug.Conn.get_req_header(conn, "upgrade") == ["websocket"]
+
+                     conn
+                   end
+                 )
+
+        assert {:ok, _client} = WebSocketClient.start_link(TestServer.url("/ws"))
+      end
     end
 
     describe "websocket_handle/3" do