Repository: bitcrowd/chromic_pdf
Pull request #64 (2020-11-09 20:51:04)

Refactor, prepare for on-demand mode.
* Extract connection stream tokenization from Connection.
* Move JsonRPC transcoding to Connection.
* Ditch CallCount process, instead move call_id to Connection.
* Merge CreateTarget & SpawnSession protocols into one.
* Pass result callback into protocol instead of calling
  GenServer.reply/2 there.


Patch:
diff --git a/lib/chromic_pdf/pdf/browser.ex b/lib/chromic_pdf/pdf/browser.ex
index 906f05a..825def1 100644
--- a/lib/chromic_pdf/pdf/browser.ex
+++ b/lib/chromic_pdf/pdf/browser.ex
@@ -2,9 +2,8 @@ defmodule ChromicPDF.Browser do
   @moduledoc false
 
   use GenServer
-  alias ChromicPDF.{CallCount, Connection, JsonRPC, Protocol}
+  alias ChromicPDF.{Connection, Protocol}
 
-  @type browser :: pid() | atom()
   @type state :: %{dispatch: Protocol.dispatch(), protocols: [Protocol.t()]}
 
   # ------------- API ----------------
@@ -22,15 +21,11 @@ defmodule ChromicPDF.Browser do
     GenServer.start_link(__MODULE__, args, name: server_name_from_args(args))
   end
 
-  @spec run_protocol(browser(), Protocol.t()) :: {:ok, any()} | {:error, term()}
+  @spec run_protocol(atom(), Protocol.t()) :: {:ok, any()} | {:error, term()}
   def run_protocol(chromic, %Protocol{} = protocol) when is_atom(chromic) do
     GenServer.call(server_name(chromic), {:run_protocol, protocol})
   end
 
-  def run_protocol(browser, %Protocol{} = protocol) when is_pid(browser) do
-    GenServer.call(browser, {:run_protocol, protocol})
-  end
-
   defp server_name_from_args(args) do
     args
     |> Keyword.fetch!(:chromic)
@@ -46,61 +41,31 @@ defmodule ChromicPDF.Browser do
   @impl GenServer
   def init(args) do
     {:ok, conn_pid} = Connection.start_link(self(), args)
-    {:ok, call_count_pid} = CallCount.start_link()
-
-    Process.flag(:trap_exit, true)
 
     dispatch = fn call ->
-      call_id = CallCount.bump(call_count_pid)
-      Connection.send_msg(conn_pid, JsonRPC.encode(call, call_id))
-      call_id
+      Connection.dispatch_call(conn_pid, call)
     end
 
     {:ok, %{dispatch: dispatch, protocols: []}}
   end
 
-  @impl GenServer
-  def terminate(:shutdown, %{dispatch: dispatch}) do
-    # Graceful shutdown: Dispatch the Browser.close message to Chrome which will cause it to
-    # detach all debugging sessions and close the port.
-    dispatch.({"Browser.close", %{}})
-
-    # The Connection process will receive a message about the port termination and forward
-    # this to us. In case Chrome takes longer than the configured supervision shutdown time,
-    # we'll receive a :brutal_kill and exit immediately, so no need for a timeout here.
-    receive do
-      {:connection_terminated, _exit_state} -> :ok
-    end
-  end
-
-  # Called when we return `{:stop, :connection_terminated, _}` from `handle_info/2`.
-  def terminate(:connection_terminated, _state), do: :ok
-
   @impl GenServer
   def handle_call(
         {:run_protocol, protocol},
         from,
         %{dispatch: dispatch, protocols: protocols} = state
       ) do
-    protocols = [Protocol.init(protocol, from, dispatch) | protocols]
-    {:noreply, update_protocols(state, protocols)}
+    protocol = Protocol.init(protocol, &GenServer.reply(from, &1), dispatch)
+    {:noreply, update_protocols(state, [protocol | protocols])}
   end
 
   # Data packets coming in from connection.
   @impl GenServer
-  def handle_info({:msg_in, data}, %{dispatch: dispatch, protocols: protocols} = state) do
-    msg = JsonRPC.decode(data)
+  def handle_info({:msg_in, msg}, %{dispatch: dispatch, protocols: protocols} = state) do
     protocols = Enum.map(protocols, &Protocol.run(&1, msg, dispatch))
     {:noreply, update_protocols(state, protocols)}
   end
 
-  def handle_info({:connection_terminated, _exit_state}, state) do
-    # If we receive this message in this `handle_info/2` clause, it means that we're not
-    # performing a graceful shutdown right now and Chrome has either crashed or was closed
-    # externally, so let's suicide and let the supervisor restart us.
-    {:stop, :connection_terminated, state}
-  end
-
   defp update_protocols(state, protocols) do
     %{state | protocols: Enum.reject(protocols, &Protocol.finished?(&1))}
   end
diff --git a/lib/chromic_pdf/pdf/call_count.ex b/lib/chromic_pdf/pdf/call_count.ex
deleted file mode 100644
index 054c7be..0000000
--- a/lib/chromic_pdf/pdf/call_count.ex
+++ /dev/null
@@ -1,13 +0,0 @@
-defmodule ChromicPDF.CallCount do
-  @moduledoc false
-
-  use Agent
-
-  def start_link(_opts \\ nil) do
-    Agent.start_link(fn -> 1 end)
-  end
-
-  def bump(pid) do
-    Agent.get_and_update(pid, &{&1, &1 + 1})
-  end
-end
diff --git a/lib/chromic_pdf/pdf/connection.ex b/lib/chromic_pdf/pdf/connection.ex
index 2f9ee64..029d20c 100644
--- a/lib/chromic_pdf/pdf/connection.ex
+++ b/lib/chromic_pdf/pdf/connection.ex
@@ -2,6 +2,7 @@ defmodule ChromicPDF.Connection do
   @moduledoc false
 
   use GenServer
+  alias ChromicPDF.Connection.{Dispatcher, JsonRPC, Tokenizer}
 
   @chrome Application.compile_env(:chromic_pdf, :chrome, ChromicPDF.ChromeImpl)
 
@@ -12,67 +13,80 @@ defmodule ChromicPDF.Connection do
     GenServer.start_link(__MODULE__, {parent_pid, opts})
   end
 
-  @spec send_msg(pid(), binary()) :: :ok
-  def send_msg(pid, msg) do
-    GenServer.cast(pid, {:send_msg, msg})
+  @spec dispatch_call(pid(), binary()) :: :ok
+  def dispatch_call(pid, msg) do
+    GenServer.call(pid, {:dispatch_call, msg})
   end
 
   # ------------ Server --------------
 
-  @impl true
+  @impl GenServer
   def init({parent_pid, opts}) do
-    chrome_opts = Keyword.take(opts, [:discard_stderr, :no_sandbox])
-    {:ok, port} = @chrome.spawn(chrome_opts)
+    {:ok, port} = spawn_chrome(opts)
+
+    Process.flag(:trap_exit, true)
 
     state = %{
       parent_pid: parent_pid,
-      port: port,
-      data: []
+      tokenizer: Tokenizer.init(),
+      dispatcher: Dispatcher.init(port)
     }
 
     {:ok, state}
   end
 
-  @impl true
-  def handle_cast({:send_msg, msg}, state) do
-    @chrome.send_msg(state.port, msg)
-    {:noreply, state}
+  defp spawn_chrome(opts) do
+    opts
+    |> Keyword.take([:discard_stderr, :no_sandbox])
+    |> @chrome.spawn()
+  end
+
+  @impl GenServer
+  def handle_call({:dispatch_call, call}, _from, state) do
+    {reply, dispatcher} = Dispatcher.dispatch(call, state.dispatcher)
+
+    {:reply, reply, %{state | dispatcher: dispatcher}}
   end
 
-  @impl true
+  @impl GenServer
   # Message from Chrome through the port.
   def handle_info({_port, {:data, data}}, state) do
-    new_state =
-      data
-      |> String.split("\0")
-      |> handle_chunks(state)
+    {msgs, tokenizer} = Tokenizer.tokenize(data, state.tokenizer)
 
-    {:noreply, new_state}
+    for msg <- msgs do
+      send(state.parent_pid, {:msg_in, JsonRPC.decode(msg)})
+    end
+
+    {:noreply, %{state | tokenizer: tokenizer}}
   end
 
   # Message triggered by Port.monitor/1.
-  def handle_info({:DOWN, _ref, :port, _port, exit_state}, state) do
-    # Notify our parent about this. We're either performing a graceful shutdown at the moment (and
-    # hence Browser is currently waiting in c:GenServer.terminate/1), or this is in fact a Chrome
-    # crash in which case the Browser can decide what to do.
-    send(state.parent_pid, {:connection_terminated, exit_state})
+  # This is unlikely to happen outside terminate/2 (:shutdown).
+  def handle_info({:DOWN, _ref, :port, _port, _exit_state}, state) do
     {:noreply, state}
   end
 
-  defp handle_chunks([blob], state), do: %{state | data: [blob | state.data]}
-  defp handle_chunks([blob, ""], state), do: handle_data(%{state | data: [blob | state.data]})
-
-  defp handle_chunks([blob | rest], state),
-    do: handle_chunks(rest, handle_data(%{state | data: [blob | state.data]}))
-
-  defp handle_data(state) do
-    msg =
-      state.data
-      |> Enum.reverse()
-      |> Enum.join()
-
-    send(state.parent_pid, {:msg_in, msg})
+  # EXIT signal from port process since we trap signals.
+  def handle_info({:EXIT, _port, _reason}, state) do
+    # Chrome has crashed or was terminated externally.
+    {:stop, :connection_terminated, state}
+  end
 
-    %{state | data: []}
+  @impl GenServer
+  def terminate(:normal, _state), do: :ok
+  def terminate(:connection_terminated, _state), do: :ok
+
+  def terminate(:shutdown, state) do
+    # Graceful shutdown: Dispatch the Browser.close call to Chrome which will cause it to detach
+    # all debugging sessions and close the port.
+    Dispatcher.dispatch({"Browser.close", %{}}, state.dispatcher)
+
+    # We can't enter the GenServer loop from here, so we need to manually receive the message
+    # about the port going down. In case Chrome takes longer than the configured supervision
+    # shutdown time, we'll receive a :brutal_kill and exit immediately, so no need for a timeout.
+    receive do
+      {:DOWN, _ref, :port, _port, _exit_state} ->
+        :ok
+    end
   end
 end
diff --git a/lib/chromic_pdf/pdf/connection/dispatcher.ex b/lib/chromic_pdf/pdf/connection/dispatcher.ex
new file mode 100644
index 0000000..3bc34b0
--- /dev/null
+++ b/lib/chromic_pdf/pdf/connection/dispatcher.ex
@@ -0,0 +1,17 @@
+defmodule ChromicPDF.Connection.Dispatcher do
+  @moduledoc false
+
+  alias ChromicPDF.Connection.JsonRPC
+
+  @chrome Application.compile_env(:chromic_pdf, :chrome, ChromicPDF.ChromeImpl)
+
+  def init(port) do
+    %{port: port, next_call_id: 1}
+  end
+
+  def dispatch(call, %{port: port, next_call_id: call_id} = state) do
+    @chrome.send_msg(port, JsonRPC.encode(call, call_id))
+
+    {call_id, %{state | next_call_id: call_id + 1}}
+  end
+end
diff --git a/lib/chromic_pdf/pdf/json_rpc.ex b/lib/chromic_pdf/pdf/connection/json_rpc.ex
similarity index 96%
rename from lib/chromic_pdf/pdf/json_rpc.ex
rename to lib/chromic_pdf/pdf/connection/json_rpc.ex
index 9a92924..248b1c0 100644
--- a/lib/chromic_pdf/pdf/json_rpc.ex
+++ b/lib/chromic_pdf/pdf/connection/json_rpc.ex
@@ -1,4 +1,4 @@
-defmodule ChromicPDF.JsonRPC do
+defmodule ChromicPDF.Connection.JsonRPC do
   @moduledoc false
 
   @type call_id :: integer()
diff --git a/lib/chromic_pdf/pdf/connection/tokenizer.ex b/lib/chromic_pdf/pdf/connection/tokenizer.ex
new file mode 100644
index 0000000..c4f885d
--- /dev/null
+++ b/lib/chromic_pdf/pdf/connection/tokenizer.ex
@@ -0,0 +1,30 @@
+defmodule ChromicPDF.Connection.Tokenizer do
+  @moduledoc false
+
+  # Returns initial memo.
+  def init do
+    []
+  end
+
+  # Returns {[msgs(), memo()]}, msgs can be consumed, memo should be saved for next data blob.
+  def tokenize(data, memo) do
+    data
+    |> String.split("\0")
+    |> handle_chunks(memo)
+  end
+
+  defp handle_chunks([blob], memo), do: {[], [blob | memo]}
+  defp handle_chunks([blob, ""], memo), do: {[join_chunks([blob | memo])], []}
+
+  defp handle_chunks([blob | rest], memo) do
+    msg = join_chunks([blob | memo])
+    {msgs, memo} = handle_chunks(rest, [])
+    {[msg | msgs], memo}
+  end
+
+  defp join_chunks(memo) do
+    memo
+    |> Enum.reverse()
+    |> Enum.join()
+  end
+end
diff --git a/lib/chromic_pdf/pdf/protocol.ex b/lib/chromic_pdf/pdf/protocol.ex
index 8bab71c..0984a9a 100644
--- a/lib/chromic_pdf/pdf/protocol.ex
+++ b/lib/chromic_pdf/pdf/protocol.ex
@@ -1,7 +1,7 @@
 defmodule ChromicPDF.Protocol do
   @moduledoc false
 
-  alias ChromicPDF.JsonRPC
+  alias ChromicPDF.Connection.JsonRPC
 
   # A protocol is a sequence of JsonRPC calls and responses/notifications.
   #
@@ -15,7 +15,7 @@ defmodule ChromicPDF.Protocol do
 
   @type state :: map()
   @type error :: {:error, term()}
-  @type step :: call_step() | await_step() | reply_step()
+  @type step :: call_step() | await_step() | output_step()
 
   @type call_fun :: (state(), dispatch() -> state() | error())
   @type call_step :: {:call, call_fun()}
@@ -23,34 +23,37 @@ defmodule ChromicPDF.Protocol do
   @type await_fun :: (state(), message() -> :no_match | {:match, state()} | error())
   @type await_step :: {:await, await_fun()}
 
-  @type reply_fun :: (state() -> any())
-  @type reply_step :: {:reply, reply_fun()}
+  @type output_fun :: (state() -> any())
+  @type output_step :: {:output, output_fun()}
+
+  @type result :: {:ok, any()} | {:error, term()}
+  @type result_fun :: (result() -> any())
 
   @type t :: %__MODULE__{
           steps: [step()],
           state: state(),
-          from: GenServer.from() | nil
+          result_fun: result_fun() | nil
         }
 
-  @enforce_keys [:steps, :state, :from]
-  defstruct [:steps, :state, :from]
+  @enforce_keys [:steps, :state, :result_fun]
+  defstruct [:steps, :state, :result_fun]
 
   @spec new([step()], state()) :: __MODULE__.t()
   def new(steps, initial_state \\ %{}) do
     %__MODULE__{
       steps: steps,
       state: initial_state,
-      from: nil
+      result_fun: nil
     }
   end
 
-  @spec init(__MODULE__.t(), GenServer.from(), dispatch()) :: __MODULE__.t()
-  def init(%__MODULE__{} = protocol, from, dispatch) do
-    advance(%{protocol | from: from}, dispatch)
+  @spec init(__MODULE__.t(), result_fun(), dispatch()) :: __MODULE__.t()
+  def init(%__MODULE__{} = protocol, result_fun, dispatch) do
+    advance(%{protocol | result_fun: result_fun}, dispatch)
   end
 
-  defp advance(%__MODULE__{state: {:error, error}, from: from} = protocol, _dispatch) do
-    GenServer.reply(from, {:error, error})
+  defp advance(%__MODULE__{state: {:error, error}, result_fun: result_fun} = protocol, _dispatch) do
+    result_fun.({:error, error})
     %{protocol | steps: []}
   end
 
@@ -63,10 +66,11 @@ defmodule ChromicPDF.Protocol do
   end
 
   defp advance(
-         %__MODULE__{steps: [{:reply, fun} | rest], state: state, from: from} = protocol,
+         %__MODULE__{steps: [{:output, output_fun} | rest], state: state, result_fun: result_fun} =
+           protocol,
          dispatch
        ) do
-    GenServer.reply(from, fun.(state))
+    result_fun.({:ok, output_fun.(state)})
     advance(%{protocol | steps: rest}, dispatch)
   end
 
diff --git a/lib/chromic_pdf/pdf/protocol_macros.ex b/lib/chromic_pdf/pdf/protocol_macros.ex
index b7baf09..28d9f7d 100644
--- a/lib/chromic_pdf/pdf/protocol_macros.ex
+++ b/lib/chromic_pdf/pdf/protocol_macros.ex
@@ -1,10 +1,11 @@
 defmodule ChromicPDF.ProtocolMacros do
   @moduledoc false
 
-  # credo:disable-for-next-line
+  # credo:disable-for-next-line Credo.Check.Refactor.CyclomaticComplexity
   defmacro steps(do: block) do
     quote do
-      alias ChromicPDF.{JsonRPC, Protocol}
+      alias ChromicPDF.Connection.JsonRPC
+      alias ChromicPDF.Protocol
 
       Module.register_attribute(__MODULE__, :steps, accumulate: true)
 
@@ -139,7 +140,8 @@ defmodule ChromicPDF.ProtocolMacros do
       def unquote(name)(state, msg) do
         last_call_id = Map.fetch!(state, :last_call_id)
 
-        if ChromicPDF.JsonRPC.is_response?(msg, last_call_id) do
+        # credo:disable-for-next-line Credo.Check.Design.AliasUsage
+        if ChromicPDF.Connection.JsonRPC.is_response?(msg, last_call_id) do
           if function_exported?(__MODULE__, unquote(cb_name), 2) do
             apply(__MODULE__, unquote(cb_name), [state, msg])
           else
@@ -164,7 +166,8 @@ defmodule ChromicPDF.ProtocolMacros do
     quote do
       @steps {:await, unquote(name), 2}
       def unquote(name)(state, msg) do
-        with true <- ChromicPDF.JsonRPC.is_notification?(msg, unquote(method)),
+        # credo:disable-for-next-line Credo.Check.Design.AliasUsage
+        with true <- ChromicPDF.Connection.JsonRPC.is_notification?(msg, unquote(method)),
              true <- state["sessionId"] == msg["sessionId"],
              true <- Enum.all?(unquote(match_keys), &notification_matches?(state, msg, &1)) do
           state = extract_from_payload(msg, "params", unquote(put_keys), state)
@@ -194,10 +197,17 @@ defmodule ChromicPDF.ProtocolMacros do
 
   def notification_matches?(state, msg, key), do: notification_matches?(state, msg, {[key], key})
 
-  defmacro reply(key) do
+  defmacro output(key) when is_binary(key) do
     quote do
-      @steps {:reply, :reply, 1}
-      def reply(state), do: {:ok, Map.fetch!(state, unquote(key))}
+      @steps {:output, :output, 1}
+      def output(state), do: Map.fetch!(state, unquote(key))
+    end
+  end
+
+  defmacro output(keys) when is_list(keys) do
+    quote do
+      @steps {:output, :output, 1}
+      def output(state), do: Map.take(state, unquote(keys))
     end
   end
 end
diff --git a/lib/chromic_pdf/pdf/protocols/capture_screenshot.ex b/lib/chromic_pdf/pdf/protocols/capture_screenshot.ex
index 628553e..2e5d3bd 100644
--- a/lib/chromic_pdf/pdf/protocols/capture_screenshot.ex
+++ b/lib/chromic_pdf/pdf/protocols/capture_screenshot.ex
@@ -19,6 +19,7 @@ defmodule ChromicPDF.CaptureScreenshot do
 
     call(:capture, "Page.captureScreenshot", &Map.get(&1, :capture_screenshot, %{}), %{})
     await_response(:captured, ["data"])
-    reply("data")
+
+    output("data")
   end
 end
diff --git a/lib/chromic_pdf/pdf/protocols/close_target.ex b/lib/chromic_pdf/pdf/protocols/close_target.ex
index c115911..17f6857 100644
--- a/lib/chromic_pdf/pdf/protocols/close_target.ex
+++ b/lib/chromic_pdf/pdf/protocols/close_target.ex
@@ -6,6 +6,6 @@ defmodule ChromicPDF.CloseTarget do
   steps do
     call(:close_target, "Target.closeTarget", [:targetId], %{})
     await_response(:target_closed, ["success"])
-    reply("success")
+    output("success")
   end
 end
diff --git a/lib/chromic_pdf/pdf/protocols/create_target.ex b/lib/chromic_pdf/pdf/protocols/create_target.ex
deleted file mode 100644
index 0f2a66d..0000000
--- a/lib/chromic_pdf/pdf/protocols/create_target.ex
+++ /dev/null
@@ -1,15 +0,0 @@
-defmodule ChromicPDF.CreateTarget do
-  @moduledoc false
-
-  import ChromicPDF.ProtocolMacros
-
-  steps do
-    call(:create_browser_context, "Target.createBrowserContext", [], %{"disposeOnDetach" => true})
-    await_response(:browser_context_created, ["browserContextId"])
-
-    call(:create_target, "Target.createTarget", ["browserContextId"], %{"url" => "about:blank"})
-    await_response(:target_created, ["targetId"])
-
-    reply("targetId")
-  end
-end
diff --git a/lib/chromic_pdf/pdf/protocols/print_to_pdf.ex b/lib/chromic_pdf/pdf/protocols/print_to_pdf.ex
index e706144..d53c466 100644
--- a/lib/chromic_pdf/pdf/protocols/print_to_pdf.ex
+++ b/lib/chromic_pdf/pdf/protocols/print_to_pdf.ex
@@ -47,6 +47,6 @@ defmodule ChromicPDF.PrintToPDF do
       await_response(:cleared, [])
     end
 
-    reply("data")
+    output("data")
   end
 end
diff --git a/lib/chromic_pdf/pdf/protocols/spawn_session.ex b/lib/chromic_pdf/pdf/protocols/spawn_session.ex
index 9bd7ac5..f9af317 100644
--- a/lib/chromic_pdf/pdf/protocols/spawn_session.ex
+++ b/lib/chromic_pdf/pdf/protocols/spawn_session.ex
@@ -6,12 +6,18 @@ defmodule ChromicPDF.SpawnSession do
   @version Mix.Project.config()[:version]
 
   steps do
-    call(:attach, "Target.attachToTarget", [:targetId], %{"flatten" => true})
+    call(:create_browser_context, "Target.createBrowserContext", [], %{"disposeOnDetach" => true})
+    await_response(:browser_context_created, ["browserContextId"])
+
+    call(:create_target, "Target.createTarget", ["browserContextId"], %{"url" => "about:blank"})
+    await_response(:target_created, ["targetId"])
+
+    call(:attach, "Target.attachToTarget", ["targetId"], %{"flatten" => true})
 
     await_notification(
       :attached,
       "Target.attachedToTarget",
-      [{["targetInfo", "targetId"], :targetId}],
+      [{["targetInfo", "targetId"], "targetId"}],
       ["sessionId"]
     )
 
@@ -35,6 +41,6 @@ defmodule ChromicPDF.SpawnSession do
 
     call(:enable_page, "Page.enable", [], %{})
 
-    reply("sessionId")
+    output(["targetId", "sessionId"])
   end
 end
diff --git a/lib/chromic_pdf/pdf/session.ex b/lib/chromic_pdf/pdf/session.ex
index d2300dc..81ea838 100644
--- a/lib/chromic_pdf/pdf/session.ex
+++ b/lib/chromic_pdf/pdf/session.ex
@@ -2,7 +2,7 @@ defmodule ChromicPDF.Session do
   @moduledoc false
 
   use GenServer
-  alias ChromicPDF.{Browser, CloseTarget, CreateTarget, SpawnSession}
+  alias ChromicPDF.{Browser, CloseTarget, SpawnSession}
 
   @default_max_session_uses 1000
 
@@ -57,11 +57,8 @@ defmodule ChromicPDF.Session do
   end
 
   defp start_session(opts) do
-    {:ok, target_id} = run_protocol_through_browser(CreateTarget.new(), opts)
-
-    {:ok, session_id} =
+    {:ok, %{"targetId" => target_id, "sessionId" => session_id}} =
       opts
-      |> Keyword.put(:targetId, target_id)
       |> Keyword.put_new(:offline, true)
       |> SpawnSession.new()
       |> run_protocol_through_browser(opts)


Test patch:
diff --git a/test/unit/chromic_pdf/pdf/call_count_test.exs b/test/unit/chromic_pdf/pdf/call_count_test.exs
deleted file mode 100644
index 984ab5a..0000000
--- a/test/unit/chromic_pdf/pdf/call_count_test.exs
+++ /dev/null
@@ -1,17 +0,0 @@
-defmodule ChromicPDF.CallCountTest do
-  use ExUnit.Case
-  alias ChromicPDF.CallCount
-
-  setup do
-    {:ok, pid} = start_supervised(CallCount)
-    %{pid: pid}
-  end
-
-  test "it starts at 1", %{pid: pid} do
-    assert CallCount.bump(pid) == 1
-  end
-
-  test "it goes up and up and up", %{pid: pid} do
-    assert Enum.map(1..5, fn _i -> CallCount.bump(pid) end) == [1, 2, 3, 4, 5]
-  end
-end
diff --git a/test/unit/chromic_pdf/pdf/connection/tokenizer_test.exs b/test/unit/chromic_pdf/pdf/connection/tokenizer_test.exs
new file mode 100644
index 0000000..cc3eeb9
--- /dev/null
+++ b/test/unit/chromic_pdf/pdf/connection/tokenizer_test.exs
@@ -0,0 +1,22 @@
+defmodule ChromicPDF.Connection.TokenizerTest do
+  use ExUnit.Case, async: true
+  import ChromicPDF.Connection.Tokenizer
+
+  describe "tokenizing incoming messages" do
+    test "single message" do
+      assert tokenize("foo\0", []) == {["foo"], []}
+    end
+
+    test "single message with previous chunks" do
+      assert tokenize("foo\0", ["bar", "baz"]) == {["bazbarfoo"], []}
+    end
+
+    test "incomplete message" do
+      assert tokenize("foo", []) == {[], ["foo"]}
+    end
+
+    test "complete message followed by incomplete message" do
+      assert tokenize("foo\0bar", []) == {["foo"], ["bar"]}
+    end
+  end
+end
diff --git a/test/unit/chromic_pdf/pdf/connection_test.exs b/test/unit/chromic_pdf/pdf/connection_test.exs
index 822a9f5..88571eb 100644
--- a/test/unit/chromic_pdf/pdf/connection_test.exs
+++ b/test/unit/chromic_pdf/pdf/connection_test.exs
@@ -2,7 +2,6 @@ defmodule ChromicPDF.ConnectionTest do
   use ExUnit.Case
   import Mox
   import ChromicPDF.Connection
-  import ChromicPDF.GenServerTestMacros
   alias ChromicPDF.ChromeMock
 
   @port :some_port
@@ -11,8 +10,11 @@ defmodule ChromicPDF.ConnectionTest do
   defp new_state do
     %{
       parent_pid: self(),
-      data: [],
-      port: @port
+      tokenizer: [],
+      dispatcher: %{
+        next_call_id: 1,
+        port: @port
+      }
     }
   end
 
@@ -20,69 +22,64 @@ defmodule ChromicPDF.ConnectionTest do
     %{state: new_state()}
   end
 
-  describe "API" do
-    test_cast ":send_msg", {:send_msg, "foo"} do
-      send_msg(self(), "foo")
-    end
-  end
-
   describe "initialization" do
     test "it spawns Chrome and initializes its state" do
-      expect(ChromeMock, :spawn, fn [no_sandbox: true] -> {:ok, @port} end)
-      assert init({self(), [no_sandbox: true]}) == {:ok, new_state()}
+      opts = [discard_stderr: false, no_sandbox: true]
+      expect(ChromeMock, :spawn, fn ^opts -> {:ok, @port} end)
+      assert init({self(), opts}) == {:ok, new_state()}
     end
   end
 
   describe "external process supervision" do
     setup [:new_state]
 
-    defp assert_connection_terminated do
-      assert_receive({:connection_terminated, 127})
-    end
+    test "it suicides when Chrome is terminated externally", %{state: state} do
+      assert handle_info({:EXIT, @port, :normal}, state) ==
+               {:stop, :connection_terminated, state}
 
-    test "it notifies the parent process when Chrome closes the pipe", %{state: state} do
-      assert handle_info({:DOWN, @ref, :port, @port, 127}, state) == {:noreply, state}
-      assert_connection_terminated()
+      assert handle_info({:EXIT, @port, :other_reason}, state) ==
+               {:stop, :connection_terminated, state}
     end
   end
 
-  describe "incoming messages" do
+  describe "graceful shutdown" do
     setup [:new_state]
 
-    defp msg_chain_out(state, msgs) do
-      Enum.reduce(msgs, state, fn msg, s ->
-        assert {:noreply, ns} = handle_info({@port, {:data, msg}}, s)
-        ns
-      end)
-    end
+    test "it gracefully closes Chrome on shutdown", %{state: %{dispatcher: %{port: port}} = state} do
+      expected_msg = ~s({"id":1,"method":"Browser.close","params":{}})
+      expect(ChromeMock, :send_msg, fn ^port, ^expected_msg -> :ok end)
 
-    defp assert_msg_in(msg) do
-      assert_receive({:msg_in, ^msg})
-    end
+      # Inject :DOWN message before calling terminate/2 to avoid locking in receive.
+      send(self(), {:DOWN, @ref, :port, @port, 0})
 
-    test "it passes received messages to its parent", %{state: state} do
-      msg_chain_out(state, ["foo\0"])
-      assert_msg_in("foo")
+      assert terminate(:shutdown, state) == :ok
     end
+  end
+
+  describe "incoming messages" do
+    setup [:new_state]
 
-    test "it can receive multiple messages in a chunk", %{state: state} do
-      msg_chain_out(state, ["foo\0bar\0"])
-      assert_msg_in("foo")
-      assert_msg_in("bar")
+    test "stores incomplete messages in the tokenizer memo", %{state: state} do
+      assert {:noreply, %{tokenizer: ["foo"]}} = handle_info({@port, {:data, "foo"}}, state)
     end
 
-    test "it can receive a long message crossing chunks", %{state: state} do
-      msg_chain_out(state, ["foo", "bar\0"])
-      assert_msg_in("foobar")
+    test "decodes complete messages and sends them to parent", %{state: state} do
+      handle_info({@port, {:data, "{}\0"}}, state)
+      assert_receive {:msg_in, %{}}
     end
   end
 
   describe "outgoing messages" do
     setup [:new_state]
 
-    test "it forwards outgoing messages to Chrome", %{state: %{port: port} = state} do
-      expect(ChromeMock, :send_msg, fn ^port, "foo" -> :ok end)
-      assert handle_cast({:send_msg, "foo"}, state) == {:noreply, state}
+    test "it encodes messages and sends them to Chrome", %{
+      state: %{dispatcher: %{port: port}} = state
+    } do
+      expected_msg = ~s({"id":1,"method":"method","params":{}})
+      expect(ChromeMock, :send_msg, fn ^port, ^expected_msg -> :ok end)
+
+      assert {:reply, 1, %{dispatcher: %{next_call_id: 2}}} =
+               handle_call({:dispatch_call, {"method", %{}}}, {self(), @ref}, state)
     end
   end
 end
diff --git a/test/unit/support/genserver_test_macros.ex b/test/unit/support/genserver_test_macros.ex
deleted file mode 100644
index 07a1cc0..0000000
--- a/test/unit/support/genserver_test_macros.ex
+++ /dev/null
@@ -1,12 +0,0 @@
-defmodule ChromicPDF.GenServerTestMacros do
-  @moduledoc false
-
-  defmacro test_cast(name, expected_cast_value, do: block) do
-    quote do
-      test "it allows to cast a #{unquote(name)} command" do
-        unquote(block)
-        assert_receive {:"$gen_cast", unquote(expected_cast_value)}
-      end
-    end
-  end
-end