feat: allow customizing expectation count via times option

lib/test_server.ex

@@ -117,7 +117,8 @@ defmodule TestServer do
   defp verify_routes!(instance) do
     instance
     |> Instance.routes()
-    |> Enum.reject(& &1.suspended)
+    |> Enum.filter(&Instance.active_route?/1)
+    |> Enum.reject(&(&1.times == :infinity))
     |> case do
       [] ->
         :ok
@@ -134,7 +135,8 @@ defmodule TestServer do
   defp verify_websocket_handlers!(instance) do
     instance
     |> Instance.websocket_handlers()
-    |> Enum.reject(& &1.suspended)
+    |> Enum.filter(&Instance.active_websocket_handler?/1)
+    |> Enum.reject(&(&1.times == :infinity))
     |> case do
       [] ->
         :ok

lib/test_server/instance.ex

@@ -193,17 +193,18 @@ defmodule TestServer.Instance do
     match = Keyword.get_lazy(options, :match, fn -> build_match_function(uri, methods) end)
 
     to = Keyword.fetch!(options, :to)
+    times = Keyword.get(options, :times, 1)
 
     route = %{
       ref: make_ref(),
       uri: uri,
       methods: methods,
       match: match,
       to: to,
-      options: Keyword.drop(options, [:via, :match, :to]),
+      options: Keyword.drop(options, [:via, :match, :to, :times]),
       requests: [],
       stacktrace: stacktrace,
-      suspended: false
+      times: times
     }
 
     {:reply, {:ok, route}, %{state | routes: state.routes ++ [route]}}
@@ -220,14 +221,16 @@ defmodule TestServer.Instance do
         _from,
         state
       ) do
+    times = Keyword.get(options, :times, 1)
+
     handler = %{
       route_ref: route_ref,
       match: Keyword.get(options, :match),
       to: Keyword.fetch!(options, :to),
-      options: Keyword.drop(options, [:match, :to]),
+      options: Keyword.drop(options, [:match, :to, :times]),
       received: [],
       stacktrace: stacktrace,
-      suspended: false
+      times: times
     }
 
     {:reply, {:ok, handler}, %{state | websocket_handlers: state.websocket_handlers ++ [handler]}}
@@ -282,6 +285,9 @@ defmodule TestServer.Instance do
     {:reply, res, state}
   end
 
+  def active_route?(route), do: route.times == :infinity or route.times > 0
+  def active_websocket_handler?(handler), do: handler.times == :infinity or handler.times > 0
+
   defp build_match_function(uri, methods) do
     {method_match, guards} =
       case methods do
@@ -359,15 +365,15 @@ defmodule TestServer.Instance do
   defp run_routes(conn, state) do
     state.routes
     |> Enum.find_index(fn
-      %{suspended: true} -> false
+      %{times: 0} -> false
       %{match: match} -> try_run_match(match, [conn])
     end)
     |> case do
       nil ->
         {{:error, {:not_found, conn}}, state}
 
       index ->
-        %{to: plug, stacktrace: stacktrace} = route = Enum.at(state.routes, index)
+        %{to: plug, stacktrace: stacktrace, times: times} = route = Enum.at(state.routes, index)
 
         result =
           conn
@@ -376,7 +382,8 @@ defmodule TestServer.Instance do
 
         routes =
           List.update_at(state.routes, index, fn route ->
-            %{route | suspended: true, requests: route.requests ++ [result]}
+            new_times = if times == :infinity, do: :infinity, else: times - 1
+            %{route | times: new_times, requests: route.requests ++ [result]}
           end)
 
         {result, %{state | routes: routes}}
@@ -405,7 +412,7 @@ defmodule TestServer.Instance do
     |> Enum.map(&{&1.route_ref == route_ref, &1})
     |> Enum.find_index(fn
       {false, _} -> false
-      {true, %{suspended: true}} -> false
+      {true, %{times: 0}} -> false
       {true, %{match: nil}} -> true
       {true, %{match: match}} -> try_run_match(match, [frame, websocket_state])
     end)
@@ -414,13 +421,14 @@ defmodule TestServer.Instance do
         {{:error, :not_found}, state}
 
       index ->
-        %{to: handler, stacktrace: stacktrace} = Enum.at(state.websocket_handlers, index)
+        %{to: handler, stacktrace: stacktrace, times: times} = Enum.at(state.websocket_handlers, index)
 
         result = try_run_websocket_handler(frame, websocket_state, stacktrace, handler)
 
         websocket_handlers =
           List.update_at(state.websocket_handlers, index, fn websocket_handle ->
-            %{websocket_handle | suspended: true, received: websocket_handle.received ++ [frame]}
+            new_times = if times == :infinity, do: :infinity, else: times - 1
+            %{websocket_handle | times: new_times, received: websocket_handle.received ++ [frame]}
           end)
 
         {result, %{state | websocket_handlers: websocket_handlers}}

lib/test_server/plug.ex

@@ -43,7 +43,7 @@ defmodule TestServer.Plug do
   end
 
   defp append_formatted_routes(message, instance) do
-    routes = Enum.split_with(Instance.routes(instance), &(not &1.suspended))
+    routes = Enum.split_with(Instance.routes(instance), &Instance.active_route?/1)
 
     """
     #{message}
@@ -52,23 +52,23 @@ defmodule TestServer.Plug do
     """
   end
 
-  defp format_routes({[], suspended_routes}) do
+  defp format_routes({[], exhausted_routes}) do
     message = "No active routes."
 
-    case suspended_routes do
+    case exhausted_routes do
       [] ->
         message
 
-      suspended_routes ->
+      exhausted_routes ->
         """
         #{message} The following routes have been processed:
 
-        #{Instance.format_routes(suspended_routes)}
+        #{Instance.format_routes(exhausted_routes)}
         """
     end
   end
 
-  defp format_routes({active_routes, _suspended_routes}) do
+  defp format_routes({active_routes, _exhausted_routes}) do
     """
     Active routes:
 

lib/test_server/websocket.ex

@@ -39,7 +39,7 @@ defmodule TestServer.WebSocket do
     websocket_handlers =
       Instance.websocket_handlers(instance)
       |> Enum.filter(&(&1.route_ref == route_ref))
-      |> Enum.split_with(&(not &1.suspended))
+      |> Enum.split_with(&Instance.active_websocket_handler?/1)
 
     """
     #{message}
@@ -48,18 +48,18 @@ defmodule TestServer.WebSocket do
     """
   end
 
-  defp format_websocket_handlers({[], suspended_websocket_handlers}) do
+  defp format_websocket_handlers({[], exhausted_websocket_handlers}) do
     message = "No active websocket handlers."
 
-    case suspended_websocket_handlers do
+    case exhausted_websocket_handlers do
       [] ->
         message
 
       websocket_handlers ->
         """
         #{message} The following websocket handlers have been processed:
 
-        #{Instance.format_routes(websocket_handlers)}"
+        #{Instance.format_websocket_handlers(websocket_handlers)}"
         """
     end
   end