From f0a8cba762bc1b3dc6bd894e715a9493bd03a928 Mon Sep 17 00:00:00 2001
From: Carlos Souza <carloshrsouza@gmail.com>
Date: Wed, 17 Jan 2024 06:26:55 -0500
Subject: [PATCH 1/4] Improve client API ergonomics.

Start moving client API functions to Xogmios top level module
to avoid direct access to Xogmios.ChainSync and Xogmios.StateQuery.
---
 lib/xogmios.ex                       | 27 ++++++++++++++++++++++-
 lib/xogmios/chain_sync.ex            | 11 +++++-----
 lib/xogmios/chain_sync/connection.ex | 11 ----------
 lib/xogmios/state_query.ex           | 32 +++++++++++++++++-----------
 test/xogmios_test.exs                |  7 +++---
 5 files changed, 56 insertions(+), 32 deletions(-)

diff --git a/lib/xogmios.ex b/lib/xogmios.ex
index 392e04d..7168a97 100644
--- a/lib/xogmios.ex
+++ b/lib/xogmios.ex
@@ -1,3 +1,28 @@
 defmodule Xogmios do
-  @moduledoc false
+  @moduledoc """
+  This is the top level module for Xogmios
+  """
+
+  alias Xogmios.ChainSync
+  alias Xogmios.StateQuery
+
+  def start_state_link(client, opts) do
+    StateQuery.start_link(client, opts)
+  end
+
+  def start_chain_sync_link(client, opts) do
+    ChainSync.start_link(client, opts)
+  end
+
+  defmacro __using__(:state_query) do
+    quote do
+      use Xogmios.StateQuery
+    end
+  end
+
+  defmacro __using__(:chain_sync) do
+    quote do
+      use Xogmios.ChainSync
+    end
+  end
 end
diff --git a/lib/xogmios/chain_sync.ex b/lib/xogmios/chain_sync.ex
index 2c61ba8..51c6d17 100644
--- a/lib/xogmios/chain_sync.ex
+++ b/lib/xogmios/chain_sync.ex
@@ -3,13 +3,17 @@ defmodule Xogmios.ChainSync do
   This module interfaces with the Chain Synchronization protocol.
   """
 
-  require Logger
-
   alias Xogmios.ChainSync.Messages
 
   @callback handle_block(map(), any()) ::
               {:ok, :next_block, map()} | {:ok, map()} | {:ok, :close, map()}
 
+  def start_link(client, opts) do
+    url = Keyword.fetch!(opts, :url)
+    state = Keyword.merge(opts, handler: client)
+    :websocket_client.start_link(url, client, state)
+  end
+
   defmacro __using__(_opts) do
     quote do
       @behaviour Xogmios.ChainSync
@@ -18,8 +22,6 @@ defmodule Xogmios.ChainSync do
 
       require Logger
 
-      @name __MODULE__
-
       def handle_message(%{"id" => "start"} = message, state) do
         %{
           "method" => "nextBlock",
@@ -64,7 +66,6 @@ defmodule Xogmios.ChainSync do
       end
 
       def handle_message({:text, message}, state) do
-        # Logger.info("fallback handle message #{inspect(message)}")
         {:close, state}
       end
     end
diff --git a/lib/xogmios/chain_sync/connection.ex b/lib/xogmios/chain_sync/connection.ex
index 5411aa1..11e0c9a 100644
--- a/lib/xogmios/chain_sync/connection.ex
+++ b/lib/xogmios/chain_sync/connection.ex
@@ -24,17 +24,6 @@ defmodule Xogmios.ChainSync.Connection do
         }
       end
 
-      def start_connection(opts),
-        do: do_start_link(opts)
-
-      def do_start_link(opts) do
-        url = Keyword.fetch!(opts, :url)
-
-        state = Keyword.merge(opts, handler: __MODULE__)
-
-        :websocket_client.start_link(url, __MODULE__, state)
-      end
-
       @impl true
       def init(state) do
         initial_state =
diff --git a/lib/xogmios/state_query.ex b/lib/xogmios/state_query.ex
index 89e98aa..7bd8577 100644
--- a/lib/xogmios/state_query.ex
+++ b/lib/xogmios/state_query.ex
@@ -8,12 +8,29 @@ defmodule Xogmios.StateQuery do
   alias Xogmios.StateQuery.Response
   alias Xogmios.StateQuery.Server
 
+  def start_link(client, opts) do
+    GenServer.start_link(client, opts, name: client)
+  end
+
   @query_messages %{
     get_current_epoch: Messages.get_current_epoch(),
     get_era_start: Messages.get_era_start()
   }
 
-  def query_messages, do: @query_messages
+  @allowed_queries Map.keys(@query_messages)
+
+  def fetch_query_message(query) when query in @allowed_queries,
+    do: Map.fetch(@query_messages, query)
+
+  def fetch_query_message(query),
+    do: {:error, "Unsupported query #{inspect(query)}"}
+
+  def call_query(client, message) do
+    case GenServer.call(client, {:send_message, message}) do
+      {:ok, response} -> {:ok, response}
+      {:error, reason} -> {:error, reason}
+    end
+  end
 
   defmacro __using__(_opts) do
     quote do
@@ -27,21 +44,12 @@ defmodule Xogmios.StateQuery do
       """
       @spec send_query(term(), term()) :: {:ok, any()} | {:error, any()}
       def send_query(client \\ __MODULE__, query) do
-        with {:ok, message} <- Map.fetch(StateQuery.query_messages(), query),
-             {:ok, %Response{} = response} <- GenServer.call(client, {:send_message, message}) do
+        with {:ok, message} <- StateQuery.fetch_query_message(query),
+             {:ok, %Response{} = response} <- StateQuery.call_query(client, message) do
           {:ok, response.result}
-        else
-          :error -> {:error, "Unsupported query"}
-          {:error, _reason} -> {:error, "Error sending query"}
         end
       end
 
-      def start_connection(opts),
-        do: do_start_link(opts)
-
-      def do_start_link(args),
-        do: GenServer.start_link(__MODULE__, args, name: __MODULE__)
-
       ## Callbacks
 
       @impl true
diff --git a/test/xogmios_test.exs b/test/xogmios_test.exs
index 8ef8ba4..aeb3be9 100644
--- a/test/xogmios_test.exs
+++ b/test/xogmios_test.exs
@@ -14,10 +14,11 @@ defmodule XogmiosTest do
   end
 
   defmodule DummyClient do
-    use Xogmios.ChainSync
+    use Xogmios, :chain_sync
 
-    def start_link(opts),
-      do: start_connection(opts)
+    def start_link(opts) do
+      Xogmios.start_chain_sync_link(__MODULE__, opts)
+    end
 
     @impl true
     def handle_block(_block, state) do

From ccd40a068540d5299d6ae7150eb81949ff41f25e Mon Sep 17 00:00:00 2001
From: Carlos Souza <carloshrsouza@gmail.com>
Date: Wed, 17 Jan 2024 09:01:08 -0500
Subject: [PATCH 2/4] Add StateQuery tests and update examples.

---
 README.md                                     | 72 +------------------
 examples/chain_sync_client.ex                 | 22 ++++--
 examples/state_query_client.ex                | 10 +--
 lib/xogmios/state_query.ex                    |  2 +-
 .../{xogmios_test.exs => chain_sync_test.exs} |  2 +-
 test/state_query_test.exs                     | 37 ++++++++++
 test/support/state_query/test_handler.ex      | 69 ++++++++++++++++++
 test/support/state_query/test_router.ex       | 25 +++++++
 test/support/state_query/test_server.ex       | 30 ++++++++
 9 files changed, 186 insertions(+), 83 deletions(-)
 rename test/{xogmios_test.exs => chain_sync_test.exs} (96%)
 create mode 100644 test/state_query_test.exs
 create mode 100644 test/support/state_query/test_handler.ex
 create mode 100644 test/support/state_query/test_router.ex
 create mode 100644 test/support/state_query/test_server.ex

diff --git a/README.md b/README.md
index 66068ef..ea82b3e 100644
--- a/README.md
+++ b/README.md
@@ -20,77 +20,7 @@ end
 
 ## Examples
 
-From a new module, call `use Xogmios.ChainSync` and implement the `start_link/1` and `handle_block/2` functions as such:
-
-```elixir
-defmodule MyApp.ChainSyncClient do
-  @moduledoc """
-  This module syncs with the tip of the chain and reads blocks indefinitely
-  """
-
-  use Xogmios.ChainSync
-
-  def start_link(opts),
-    do: start_connection(opts)
-
-  @impl true
-  def handle_block(block, state) do
-    IO.puts("handle_block #{block["height"]}")
-    {:ok, :next_block, state}
-  end
-end
-```
-
-Add the new module to your app's supervision tree in `application.ex`:
-
-```elixir
-def start(_type, _args) do
-  children =[
-    {MyApp.ChainSyncClient, url: "ws://url-for-ogmios"},
-  ]
-
-  opts = [strategy: :one_for_one, name: MyApp.Supervisor]
-  Supervisor.start_link(children, opts)
-end
-```
-
-The example below syncs with the tip and prints the next 3 blocks:
-
-```elixir
-defmodule MyApp.ChainSyncClient do
-  @moduledoc """
-  This module syncs with the tip of the chain and reads the following 3 blocks
-  """
-
-  use Xogmios.ChainSync
-
-  require Logger
-
-  def start_link(opts) do
-    # Initial state currently has to be defined here
-    # and passed as argument to start_connection
-    initial_state = [counter: 3]
-
-    opts
-    |> Keyword.merge(initial_state)
-    |> start_connection()
-  end
-
-  @impl true
-  def handle_block(block, %{counter: counter} = state) when counter > 1 do
-    Logger.info("handle_block #{block["height"]}")
-    {:ok, :next_block, %{state | counter: counter - 1}}
-  end
-
-  @impl true
-  def handle_block(block, state) do
-    Logger.info("final handle_block #{block["height"]}")
-    {:ok, :close, state}
-  end
-end
-```
-
-See more examples in the [examples](./examples/) folder.
+See [ChainSyncClient](./examples/chain_sync_client.ex) and [StateQueryClient](./examples/state_query_client.ex)
 
 ## Test
 
diff --git a/examples/chain_sync_client.ex b/examples/chain_sync_client.ex
index 38fc210..41fdcd1 100644
--- a/examples/chain_sync_client.ex
+++ b/examples/chain_sync_client.ex
@@ -1,6 +1,7 @@
 defmodule ChainSyncClient do
   @moduledoc """
-  This module syncs with the tip of the chain and reads blocks indefinitely.
+  This module syncs with the tip of the chain, reads the next 3 blocks
+  and then closes the connection with the server.
 
   Add this to your application's supervision tree like so:
 
@@ -12,14 +13,23 @@ defmodule ChainSyncClient do
   end
   """
 
-  use Xogmios.ChainSync
+  use Xogmios, :chain_sync
 
-  def start_link(opts),
-    do: start_connection(opts)
+  def start_link(opts) do
+    initial_state = [counter: 3]
+    opts = Keyword.merge(opts, initial_state)
+    Xogmios.start_chain_sync_link(__MODULE__, opts)
+  end
+
+  @impl true
+  def handle_block(block, %{counter: counter} = state) when counter > 1 do
+    IO.puts("handle_block #{block["height"]}")
+    {:ok, :next_block, %{state | counter: counter - 1}}
+  end
 
   @impl true
   def handle_block(block, state) do
-    IO.puts("handle_block from client #{block["height"]}")
-    {:ok, :next_block, state}
+    IO.puts("final handle_block #{block["height"]}")
+    {:ok, :close, state}
   end
 end
diff --git a/examples/state_query_client.ex b/examples/state_query_client.ex
index 30b9ca8..95f1f25 100644
--- a/examples/state_query_client.ex
+++ b/examples/state_query_client.ex
@@ -13,14 +13,16 @@ defmodule StateQueryClient do
 
   Then invoke functions:
    * StateQueryClient.get_current_epoch()
-   * StateQueryClient.get_era_start()
+   * StateQueryClient.get_era_start() # not yet supported
 
   Not all queries are supported yet.
   """
-  use Xogmios.StateQuery
 
-  def start_link(opts),
-    do: start_connection(opts)
+  use Xogmios, :state_query
+
+  def start_link(opts) do
+    Xogmios.start_state_link(__MODULE__, opts)
+  end
 
   def get_current_epoch() do
     case send_query(:get_current_epoch) do
diff --git a/lib/xogmios/state_query.ex b/lib/xogmios/state_query.ex
index 7bd8577..97d60f6 100644
--- a/lib/xogmios/state_query.ex
+++ b/lib/xogmios/state_query.ex
@@ -67,7 +67,7 @@ defmodule Xogmios.StateQuery do
 
       @impl true
       def handle_call({:send_message, message}, from, state) do
-        send(state.ws_pid, {:store_caller, from})
+        {:store_caller, _from} = send(state.ws_pid, {:store_caller, from})
         :ok = :websocket_client.send(state.ws_pid, {:text, message})
         {:noreply, state}
       end
diff --git a/test/xogmios_test.exs b/test/chain_sync_test.exs
similarity index 96%
rename from test/xogmios_test.exs
rename to test/chain_sync_test.exs
index aeb3be9..3084ba5 100644
--- a/test/xogmios_test.exs
+++ b/test/chain_sync_test.exs
@@ -1,4 +1,4 @@
-defmodule XogmiosTest do
+defmodule Xogmios.ChainSyncTest do
   use ExUnit.Case
 
   @ws_url ChainSync.TestServer.get_url()
diff --git a/test/state_query_test.exs b/test/state_query_test.exs
new file mode 100644
index 0000000..5cc176a
--- /dev/null
+++ b/test/state_query_test.exs
@@ -0,0 +1,37 @@
+defmodule Xogmios.StateQueryTest do
+  use ExUnit.Case
+
+  @ws_url StateQuery.TestServer.get_url()
+
+  setup_all do
+    {:ok, _server} = StateQuery.TestServer.start()
+
+    on_exit(fn ->
+      StateQuery.TestServer.shutdown()
+    end)
+
+    :ok
+  end
+
+  defmodule DummyClient do
+    use Xogmios, :state_query
+
+    def start_link(opts) do
+      Xogmios.start_state_link(__MODULE__, opts)
+    end
+
+    def get_current_epoch() do
+      case send_query(:get_current_epoch) do
+        {:ok, result} -> result
+        {:error, reason} -> "Something went wrong #{inspect(reason)}"
+      end
+    end
+  end
+
+  test "returns current epoch" do
+    pid = start_supervised!({DummyClient, url: @ws_url})
+    assert is_pid(pid)
+    Process.sleep(1_000)
+    assert DummyClient.get_current_epoch() == 333
+  end
+end
diff --git a/test/support/state_query/test_handler.ex b/test/support/state_query/test_handler.ex
new file mode 100644
index 0000000..8997e4d
--- /dev/null
+++ b/test/support/state_query/test_handler.ex
@@ -0,0 +1,69 @@
+defmodule StateQuery.TestHandler do
+  @moduledoc false
+
+  require Logger
+
+  @behaviour :cowboy_websocket
+  @current_epoch 333
+
+  @impl true
+  def init(request, state) do
+    {:cowboy_websocket, request, state}
+  end
+
+  @impl true
+  def websocket_init(state) do
+    {:ok, state}
+  end
+
+  @impl true
+  # Sends response back to client
+  def websocket_handle({:text, payload}, state) do
+    case Jason.decode(payload) do
+      {:ok, %{"method" => "queryNetwork/tip"}} ->
+        payload =
+          Jason.encode!(%{
+            "method" => "queryNetwork/tip",
+            "result" => %{"slot" => "123", "id" => "123"}
+          })
+
+        {:reply, {:text, payload}, state}
+
+      {:ok, %{"method" => "acquireLedgerState"}} ->
+        payload =
+          Jason.encode!(%{
+            "method" => "acquireLedgerState"
+          })
+
+        {:reply, {:text, payload}, state}
+
+      {:ok, %{"method" => "queryLedgerState/epoch"}} ->
+        payload =
+          Jason.encode!(%{
+            "method" => "can-be-anything",
+            "result" => @current_epoch
+          })
+
+        {:reply, {:text, payload}, state}
+
+      result ->
+        IO.puts("Did not match #{inspect(result)}")
+        {:reply, {:text, payload}, state}
+    end
+  end
+
+  @impl true
+  def terminate(_arg1, _arg2, _arg3) do
+    :ok
+  end
+
+  @impl true
+  def websocket_info(:stop, state) do
+    {:stop, state}
+  end
+
+  @impl true
+  def websocket_info(_message, state) do
+    {:stop, state}
+  end
+end
diff --git a/test/support/state_query/test_router.ex b/test/support/state_query/test_router.ex
new file mode 100644
index 0000000..11e8ba1
--- /dev/null
+++ b/test/support/state_query/test_router.ex
@@ -0,0 +1,25 @@
+defmodule StateQuery.TestRouter do
+  @moduledoc false
+
+  use Plug.Router
+
+  plug(:match)
+  plug(:dispatch)
+
+  match _ do
+    send_resp(conn, 200, "waiting for ws")
+  end
+
+  def options(opts) do
+    dispatch = [
+      {:_,
+       [
+         {"/ws/[...]", StateQuery.TestHandler, []}
+       ]}
+    ]
+
+    port = opts[:port] || 8088
+
+    [port: port, dispatch: dispatch]
+  end
+end
diff --git a/test/support/state_query/test_server.ex b/test/support/state_query/test_server.ex
new file mode 100644
index 0000000..e3f9a2f
--- /dev/null
+++ b/test/support/state_query/test_server.ex
@@ -0,0 +1,30 @@
+defmodule StateQuery.TestServer do
+  @moduledoc false
+
+  @default_port 8989
+
+  def get_url(port \\ @default_port) do
+    "ws://localhost:#{port}/ws"
+  end
+
+  def start(port \\ @default_port) do
+    cowboy_server =
+      Plug.Cowboy.http(
+        WebSocket.Router,
+        [scheme: :http],
+        StateQuery.TestRouter.options(port: port)
+      )
+
+    cowboy_server =
+      case cowboy_server do
+        {:error, {:already_started, server}} -> server
+        pid -> pid
+      end
+
+    {:ok, cowboy_server}
+  end
+
+  def shutdown() do
+    Plug.Cowboy.shutdown(WebSocket.Router.HTTP)
+  end
+end

From cc140271862ffd0618231b8b32b8b6015933f722 Mon Sep 17 00:00:00 2001
From: Carlos Souza <carloshrsouza@gmail.com>
Date: Wed, 17 Jan 2024 09:06:44 -0500
Subject: [PATCH 3/4] Add test for invalid query

---
 test/state_query_test.exs | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/test/state_query_test.exs b/test/state_query_test.exs
index 5cc176a..b354273 100644
--- a/test/state_query_test.exs
+++ b/test/state_query_test.exs
@@ -26,6 +26,13 @@ defmodule Xogmios.StateQueryTest do
         {:error, reason} -> "Something went wrong #{inspect(reason)}"
       end
     end
+
+    def get_bananas() do
+      case send_query(:get_bananas) do
+        {:ok, result} -> result
+        {:error, reason} -> "Something went wrong #{inspect(reason)}"
+      end
+    end
   end
 
   test "returns current epoch" do
@@ -33,5 +40,6 @@ defmodule Xogmios.StateQueryTest do
     assert is_pid(pid)
     Process.sleep(1_000)
     assert DummyClient.get_current_epoch() == 333
+    assert DummyClient.get_bananas() =~ "Something went wrong"
   end
 end

From d7574cd727b4fdd76e9a5f1dc828a3cd4699df7c Mon Sep 17 00:00:00 2001
From: Carlos Souza <carloshrsouza@gmail.com>
Date: Wed, 17 Jan 2024 09:14:12 -0500
Subject: [PATCH 4/4] Clean up tests

---
 test/chain_sync_test.exs                     |  6 ++--
 test/state_query_test.exs                    |  6 ++--
 test/support/state_query/test_router.ex      | 25 ----------------
 test/support/state_query/test_server.ex      | 30 --------------------
 test/support/{chain_sync => }/test_router.ex |  4 +--
 test/support/{chain_sync => }/test_server.ex |  6 ++--
 6 files changed, 11 insertions(+), 66 deletions(-)
 delete mode 100644 test/support/state_query/test_router.ex
 delete mode 100644 test/support/state_query/test_server.ex
 rename test/support/{chain_sync => }/test_router.ex (77%)
 rename test/support/{chain_sync => }/test_server.ex (77%)

diff --git a/test/chain_sync_test.exs b/test/chain_sync_test.exs
index 3084ba5..af00558 100644
--- a/test/chain_sync_test.exs
+++ b/test/chain_sync_test.exs
@@ -1,13 +1,13 @@
 defmodule Xogmios.ChainSyncTest do
   use ExUnit.Case
 
-  @ws_url ChainSync.TestServer.get_url()
+  @ws_url TestServer.get_url()
 
   setup_all do
-    {:ok, _server} = ChainSync.TestServer.start()
+    {:ok, _server} = TestServer.start(handler: ChainSync.TestHandler)
 
     on_exit(fn ->
-      ChainSync.TestServer.shutdown()
+      TestServer.shutdown()
     end)
 
     :ok
diff --git a/test/state_query_test.exs b/test/state_query_test.exs
index b354273..6889134 100644
--- a/test/state_query_test.exs
+++ b/test/state_query_test.exs
@@ -1,13 +1,13 @@
 defmodule Xogmios.StateQueryTest do
   use ExUnit.Case
 
-  @ws_url StateQuery.TestServer.get_url()
+  @ws_url TestServer.get_url()
 
   setup_all do
-    {:ok, _server} = StateQuery.TestServer.start()
+    {:ok, _server} = TestServer.start(handler: StateQuery.TestHandler)
 
     on_exit(fn ->
-      StateQuery.TestServer.shutdown()
+      TestServer.shutdown()
     end)
 
     :ok
diff --git a/test/support/state_query/test_router.ex b/test/support/state_query/test_router.ex
deleted file mode 100644
index 11e8ba1..0000000
--- a/test/support/state_query/test_router.ex
+++ /dev/null
@@ -1,25 +0,0 @@
-defmodule StateQuery.TestRouter do
-  @moduledoc false
-
-  use Plug.Router
-
-  plug(:match)
-  plug(:dispatch)
-
-  match _ do
-    send_resp(conn, 200, "waiting for ws")
-  end
-
-  def options(opts) do
-    dispatch = [
-      {:_,
-       [
-         {"/ws/[...]", StateQuery.TestHandler, []}
-       ]}
-    ]
-
-    port = opts[:port] || 8088
-
-    [port: port, dispatch: dispatch]
-  end
-end
diff --git a/test/support/state_query/test_server.ex b/test/support/state_query/test_server.ex
deleted file mode 100644
index e3f9a2f..0000000
--- a/test/support/state_query/test_server.ex
+++ /dev/null
@@ -1,30 +0,0 @@
-defmodule StateQuery.TestServer do
-  @moduledoc false
-
-  @default_port 8989
-
-  def get_url(port \\ @default_port) do
-    "ws://localhost:#{port}/ws"
-  end
-
-  def start(port \\ @default_port) do
-    cowboy_server =
-      Plug.Cowboy.http(
-        WebSocket.Router,
-        [scheme: :http],
-        StateQuery.TestRouter.options(port: port)
-      )
-
-    cowboy_server =
-      case cowboy_server do
-        {:error, {:already_started, server}} -> server
-        pid -> pid
-      end
-
-    {:ok, cowboy_server}
-  end
-
-  def shutdown() do
-    Plug.Cowboy.shutdown(WebSocket.Router.HTTP)
-  end
-end
diff --git a/test/support/chain_sync/test_router.ex b/test/support/test_router.ex
similarity index 77%
rename from test/support/chain_sync/test_router.ex
rename to test/support/test_router.ex
index 2900ec7..3a5307e 100644
--- a/test/support/chain_sync/test_router.ex
+++ b/test/support/test_router.ex
@@ -1,4 +1,4 @@
-defmodule ChainSync.TestRouter do
+defmodule TestRouter do
   @moduledoc false
 
   use Plug.Router
@@ -14,7 +14,7 @@ defmodule ChainSync.TestRouter do
     dispatch = [
       {:_,
        [
-         {"/ws/[...]", ChainSync.TestHandler, []}
+         {"/ws/[...]", opts[:handler], []}
        ]}
     ]
 
diff --git a/test/support/chain_sync/test_server.ex b/test/support/test_server.ex
similarity index 77%
rename from test/support/chain_sync/test_server.ex
rename to test/support/test_server.ex
index 81940c6..eac572f 100644
--- a/test/support/chain_sync/test_server.ex
+++ b/test/support/test_server.ex
@@ -1,4 +1,4 @@
-defmodule ChainSync.TestServer do
+defmodule TestServer do
   @moduledoc false
 
   @default_port 8989
@@ -7,12 +7,12 @@ defmodule ChainSync.TestServer do
     "ws://localhost:#{port}/ws"
   end
 
-  def start(port \\ @default_port) do
+  def start(port \\ @default_port, handler: handler) do
     cowboy_server =
       Plug.Cowboy.http(
         WebSocket.Router,
         [scheme: :http],
-        ChainSync.TestRouter.options(port: port)
+        TestRouter.options(port: port, handler: handler)
       )
 
     cowboy_server =