Take callback as second argument

CHANGELOG.md

@@ -1,3 +1,7 @@
+## v0.1.5 (TBA)
+
+- `TestServer.websocket_info/2` now takes the callback function as second argument
+
 ## v0.1.4 (2022-09-13)
 
 - WebSocket support

README.md

@@ -122,7 +122,7 @@ TestServer.websocket_handle(socket)
 TestServer.websocket_handle(socket, to: fn _frame_, state -> {:reply, "pong", state})
 TestServer.websocket_handle(socket, match: fn {_opcode, message}, _state -> messsage == "ping")
 
-TestServer.websocket_info(socket, to: fn state -> {:reply, {:text, "ping"}, state} end)
+TestServer.websocket_info(socket, fn state -> {:reply, {:text, "ping"}, state} end)
 ```
 
 <!-- MDOC !-->

lib/test_server.ex

@@ -441,16 +441,17 @@ defmodule TestServer do
   @doc """
   Sends an message to a websocket instance.
   """
-  @spec websocket_info(websocket_socket(), keyword()) :: :ok
-  def websocket_info({instance, _route_ref} = socket, options \\ []) do
+  @spec websocket_info(websocket_socket(), function() | nil) :: :ok
+  def websocket_info({instance, _route_ref} = socket, callback \\ nil)
+      when is_function(callback) or is_nil(callback) do
     instance_alive!(instance)
 
     [_first_module_entry | stacktrace] = get_stacktrace()
 
-    options = Keyword.put_new(options, :to, &default_websocket_info/1)
+    callback = callback || (&default_websocket_info/1)
 
     for pid <- Instance.active_websocket_connections(socket) do
-      send(pid, {options, stacktrace})
+      send(pid, {callback, stacktrace})
     end
 
     :ok

lib/test_server/instance.ex

@@ -33,7 +33,10 @@ defmodule TestServer.Instance do
     GenServer.call(instance, {:register, {:plug, {plug, stacktrace}}})
   end
 
-  @spec register(TestServer.websocket_socket(), {:websocket, {:handle, keyword(), TestServer.stacktrace()}}) ::
+  @spec register(
+          TestServer.websocket_socket(),
+          {:websocket, {:handle, keyword(), TestServer.stacktrace()}}
+        ) ::
           {:ok, map()}
   def register({instance, _router_ref} = socket, {:websocket, {:handle, options, stacktrace}}) do
     unless is_nil(options[:match]) or is_function(options[:match]),
@@ -66,17 +69,17 @@ defmodule TestServer.Instance do
 
   @spec dispatch(
           TestServer.websocket_socket(),
-          {:websocket, {:info, keyword(), TestServer.stacktrace()}, TestServer.websocket_state()}
+          {:websocket, {:info, function(), TestServer.stacktrace()}, TestServer.websocket_state()}
         ) ::
           {:ok, TestServer.websocket_reply()}
           | {:error, {term(), TestServer.stacktrace()}}
   def dispatch(
         {instance, _router_ref} = socket,
-        {:websocket, {:info, options, stacktrace}, state}
+        {:websocket, {:info, callback, stacktrace}, state}
       ) do
     GenServer.call(
       instance,
-      {:dispatch, {:websocket, socket, {:info, options, stacktrace}, state}}
+      {:dispatch, {:websocket, socket, {:info, callback, stacktrace}, state}}
     )
   end
 
@@ -286,16 +289,13 @@ defmodule TestServer.Instance do
   end
 
   def handle_call(
-        {:dispatch, {:websocket, _socket, {:info, options, stacktrace}, websocket_state}},
+        {:dispatch, {:websocket, _socket, {:info, callback, stacktrace}, websocket_state}},
         _from,
         state
       ) do
     res =
       try do
-        unless is_function(options[:to]),
-          do: raise(BadFunctionError, term: options[:to])
-
-        frame = validate_websocket_frame!(options[:to].(websocket_state), stacktrace)
+        frame = validate_websocket_frame!(callback.(websocket_state), stacktrace)
 
         {:ok, frame}
       rescue

lib/test_server/plug_cowboy/handler.ex

@@ -80,8 +80,8 @@ defmodule TestServer.Plug.Cowboy.Handler do
   defp handle_reply({:ok, state}, socket), do: {:ok, {socket, state}}
 
   @impl true
-  def websocket_info({options, stacktrace}, {socket, state}) do
-    case Instance.dispatch(socket, {:websocket, {:info, options, stacktrace}, state}) do
+  def websocket_info({callback, stacktrace}, {socket, state}) do
+    case Instance.dispatch(socket, {:websocket, {:info, callback, stacktrace}, state}) do
       {:ok, result} -> handle_reply(result, socket)
       {:error, {error, stacktrace}} -> reply_with_error({socket, state}, {error, stacktrace})
     end

test/test_server_test.exs

@@ -598,12 +598,6 @@ defmodule TestServerTest do
       end
     end
 
-    test "invalid options" do
-      assert_raise ArgumentError, ~r/`:to` is an invalid option/, fn ->
-        TestServer.websocket_init("/", to: :invalid)
-      end
-    end
-
     test "with invalid callback response" do
       defmodule WebSocketInfoInvalidMessageTest do
         use ExUnit.Case
@@ -612,7 +606,7 @@ defmodule TestServerTest do
           {:ok, _instance} = TestServer.start(suppress_warning: true)
           assert {:ok, socket} = TestServer.websocket_init("/ws")
           assert {:ok, client} = WebSocketClient.start_link(TestServer.url("/ws"))
-          assert :ok = TestServer.websocket_info(socket, to: fn _state -> :invalid end)
+          assert :ok = TestServer.websocket_info(socket, fn _state -> :invalid end)
 
           assert {:ok, message} = WebSocketClient.receive_message(client)
           assert message =~ "(RuntimeError) Invalid callback response, got: :invalid."
@@ -632,7 +626,7 @@ defmodule TestServerTest do
           assert {:ok, socket} = TestServer.websocket_init("/ws")
           assert {:ok, client} = WebSocketClient.start_link(TestServer.url("/ws"))
 
-          assert :ok = TestServer.websocket_info(socket, to: fn _state -> raise "boom" end)
+          assert :ok = TestServer.websocket_info(socket, fn _state -> raise "boom" end)
 
           assert {:ok, message} = WebSocketClient.receive_message(client)
           assert message =~ "(RuntimeError) boom"
@@ -649,11 +643,9 @@ defmodule TestServerTest do
       assert {:ok, client} = WebSocketClient.start_link(TestServer.url("/ws"))
 
       assert :ok =
-               TestServer.websocket_info(socket,
-                 to: fn state ->
-                   {:reply, {:text, "pong"}, state}
-                 end
-               )
+               TestServer.websocket_info(socket, fn state ->
+                 {:reply, {:text, "pong"}, state}
+               end)
 
       assert {:ok, "pong"} = WebSocketClient.receive_message(client)
     end