dexcord/test/support/fake_gateway.ex

defmodule Dexcord.FakeGateway do
  @moduledoc """
  A scripted, in-process Discord gateway for integration tests.

  A `Bandit` server on an ephemeral port speaks the real websocket protocol via a
  thin `WebSock` handler (`Dexcord.FakeGateway.Handler`); all decision-making lives
  in this GenServer. Every frame a client sends is decoded and forwarded to the
  controlling test process as `{:fake_gw, :frame, conn_id, map}`; connection
  lifecycle is surfaced as `{:fake_gw, :connect, conn_id}` /
  `{:fake_gw, :disconnect, conn_id}`. `conn_id` is a monotonically increasing
  1-based index so a test can tell a reconnection's frames apart from the first
  connection's.

  ## Automatic reactions (configurable per test)

    * `on_connect`  - `:hello` (send op 10 with `hello_interval`, default),
      `:silent` (wait for a manual `push_hello/2`), or `{:close, code}`.
    * `on_identify` - reaction to a client op 2: `{:ready, session_id}` (default,
      sends a READY dispatch whose `resume_gateway_url` points back at this fake
      server), `:manual`, `{:close, code}`, or `{:invalid, d}` (op 9).
    * `on_resume`   - reaction to a client op 6: `:resumed` (default),
      `{:replay, events, :resumed}` (each `{name, data, seq}` as an op 0 dispatch,
      then RESUMED), `:manual`, `{:close, code}`, or `{:invalid, d}` (op 9).
    * `auto_ack`    - reply op 11 to client heartbeats (default `true`; disable to
      simulate a zombie connection).

  Manual directives (`push_*`, `close/2`, `tcp_kill/1`) are routed to the most
  recent connection. `tcp_kill/1` brutally kills that connection process, dropping
  the TCP socket with no close frame - a laptop-sleep / cable-pull simulation.
  """

  use GenServer

  # --- Public API ---------------------------------------------------------

  def start_link(opts \\ []) do
    GenServer.start_link(__MODULE__, opts)
  end

  def child_spec(opts) do
    %{id: Keyword.get(opts, :id, __MODULE__), start: {__MODULE__, :start_link, [opts]}}
  end

  @doc "The bound TCP port."
  def port(server), do: GenServer.call(server, :port)

  @doc "The `ws://` URL clients should connect to."
  def url(server), do: GenServer.call(server, :url)

  @doc "Total number of connections accepted so far."
  def conn_count(server), do: GenServer.call(server, :conn_count)

  @doc "All decoded client frames received, in order, as `{conn_id, map}`."
  def frames(server), do: GenServer.call(server, :frames)

  @doc "Sets the reaction to a client IDENTIFY (op 2)."
  def on_identify(server, behavior), do: GenServer.call(server, {:set, :on_identify, behavior})

  @doc "Sets the reaction to a client RESUME (op 6)."
  def on_resume(server, behavior), do: GenServer.call(server, {:set, :on_resume, behavior})

  @doc "Sets the reaction to a new connection."
  def on_connect(server, behavior), do: GenServer.call(server, {:set, :on_connect, behavior})

  @doc "Toggles automatic heartbeat ACKs (op 11)."
  def auto_ack(server, bool?), do: GenServer.call(server, {:set, :auto_ack, bool?})

  @doc "Manually push a HELLO (op 10) on the current connection."
  def push_hello(server, interval \\ nil), do: GenServer.call(server, {:push_hello, interval})

  @doc "Manually push a READY dispatch."
  def push_ready(server, session_id \\ nil, resume_url \\ nil),
    do: GenServer.call(server, {:push_ready, session_id, resume_url})

  @doc "Push a RESUMED dispatch."
  def push_resumed(server, seq \\ nil), do: GenServer.call(server, {:push_resumed, seq})

  @doc "Push an op 0 dispatch with the given event name, data and sequence number."
  def push_dispatch(server, name, data, seq),
    do: GenServer.call(server, {:push_dispatch, name, data, seq})

  @doc "Push a server-requested heartbeat (op 1)."
  def push_op1(server), do: GenServer.call(server, {:push_op, 1})

  @doc "Push a RECONNECT (op 7)."
  def push_op7(server), do: GenServer.call(server, {:push_op, 7})

  @doc "Push an INVALID SESSION (op 9) with the given `d` boolean."
  def push_invalid(server, d) when is_boolean(d),
    do: GenServer.call(server, {:push_invalid, d})

  @doc "Push a heartbeat ACK (op 11)."
  def push_ack(server), do: GenServer.call(server, {:push_op, 11})

  @doc "Send a close frame with an arbitrary code on the current connection."
  def close(server, code), do: GenServer.call(server, {:close, code})

  @doc "Brutally kill the current connection process (no close frame)."
  def tcp_kill(server), do: GenServer.call(server, :tcp_kill)

  @doc false
  # Called synchronously by a handler from its `init/1`. Registers the connection
  # and returns the initial action so HELLO (or an immediate close) is emitted as
  # the handler's very first frame - no round-trip race.
  def register(server, handler_pid), do: GenServer.call(server, {:register, handler_pid})

  # --- GenServer ----------------------------------------------------------

  @impl true
  def init(opts) do
    test_pid = Keyword.get(opts, :test_pid, self())

    {:ok, bandit} =
      Bandit.start_link(
        plug: {Dexcord.FakeGateway.Plug, server: self()},
        scheme: :http,
        ip: {127, 0, 0, 1},
        port: 0,
        startup_log: false
      )

    {:ok, {_addr, port}} = ThousandIsland.listener_info(bandit)

    state = %{
      test_pid: test_pid,
      bandit: bandit,
      port: port,
      session_id: Keyword.get(opts, :session_id, "fake-session"),
      hello_interval: Keyword.get(opts, :hello_interval, 100),
      on_connect: Keyword.get(opts, :on_connect, :hello),
      on_identify: Keyword.get(opts, :on_identify, {:ready, nil}),
      on_resume: Keyword.get(opts, :on_resume, :resumed),
      auto_ack: Keyword.get(opts, :auto_ack, true),
      conn_count: 0,
      conns: %{},
      current: nil,
      frames: []
    }

    {:ok, state}
  end

  @impl true
  def handle_call(:port, _from, state), do: {:reply, state.port, state}
  def handle_call(:url, _from, state), do: {:reply, ws_url(state), state}
  def handle_call(:conn_count, _from, state), do: {:reply, state.conn_count, state}
  def handle_call(:frames, _from, state), do: {:reply, Enum.reverse(state.frames), state}

  def handle_call({:set, key, value}, _from, state),
    do: {:reply, :ok, Map.put(state, key, value)}

  def handle_call({:push_hello, interval}, _from, state) do
    push(state, hello_frame(interval || state.hello_interval))
    {:reply, :ok, state}
  end

  def handle_call({:push_ready, sid, resume_url}, _from, state) do
    push(state, ready_frame(sid || state.session_id, resume_url || ws_url(state)))
    {:reply, :ok, state}
  end

  def handle_call({:push_resumed, seq}, _from, state) do
    push(state, resumed_frame(seq))
    {:reply, :ok, state}
  end

  def handle_call({:push_dispatch, name, data, seq}, _from, state) do
    push(state, dispatch_frame(name, data, seq))
    {:reply, :ok, state}
  end

  def handle_call({:push_op, op}, _from, state) do
    push(state, %{"op" => op})
    {:reply, :ok, state}
  end

  def handle_call({:push_invalid, d}, _from, state) do
    push(state, %{"op" => 9, "d" => d})
    {:reply, :ok, state}
  end

  def handle_call({:close, code}, _from, state) do
    if state.current, do: send(state.current, {:close, code})
    {:reply, :ok, state}
  end

  def handle_call(:tcp_kill, _from, state) do
    if state.current, do: Process.exit(state.current, :kill)
    {:reply, :ok, state}
  end

  def handle_call({:register, pid}, _from, state) do
    conn_id = state.conn_count + 1
    Process.monitor(pid)

    state = %{
      state
      | conn_count: conn_id,
        conns: Map.put(state.conns, pid, conn_id),
        current: pid
    }

    send(state.test_pid, {:fake_gw, :connect, conn_id})
    {:reply, {conn_id, on_connect_action(state)}, state}
  end

  @impl true
  def handle_info({:client_frame, pid, raw}, state) do
    map = JSON.decode!(raw)
    conn_id = Map.get(state.conns, pid)
    send(state.test_pid, {:fake_gw, :frame, conn_id, map})
    state = %{state | frames: [{conn_id, map} | state.frames]}
    react(state, pid, map)
    {:noreply, state}
  end

  def handle_info({:DOWN, _ref, :process, pid, _reason}, state) do
    conn_id = Map.get(state.conns, pid)
    if conn_id, do: send(state.test_pid, {:fake_gw, :disconnect, conn_id})
    current = if state.current == pid, do: nil, else: state.current
    {:noreply, %{state | conns: Map.delete(state.conns, pid), current: current}}
  end

  def handle_info(_msg, state), do: {:noreply, state}

  # --- automatic reactions ------------------------------------------------

  # Initial action returned to a handler's init so it emits HELLO (or a close) as
  # its first frame, with no window for a client frame to slip in first.
  defp on_connect_action(state) do
    case state.on_connect do
      :hello -> {:push, encode(hello_frame(state.hello_interval))}
      :silent -> :none
      {:close, code} -> {:close, code}
    end
  end

  # op 1: client heartbeat -> ACK unless zombie mode.
  defp react(state, pid, %{"op" => 1}) do
    if state.auto_ack, do: send(pid, {:push, encode(%{"op" => 11})})
  end

  # op 2: client IDENTIFY.
  defp react(state, pid, %{"op" => 2}) do
    case state.on_identify do
      {:ready, sid} ->
        send(pid, {:push, encode(ready_frame(sid || state.session_id, ws_url(state)))})

      :manual ->
        :ok

      {:close, code} ->
        send(pid, {:close, code})

      {:invalid, d} ->
        send(pid, {:push, encode(%{"op" => 9, "d" => d})})
    end
  end

  # op 6: client RESUME.
  defp react(state, pid, %{"op" => 6}) do
    case state.on_resume do
      :resumed ->
        send(pid, {:push, encode(resumed_frame(nil))})

      {:replay, events, then} ->
        for {name, data, seq} <- events do
          send(pid, {:push, encode(dispatch_frame(name, data, seq))})
        end

        if then == :resumed, do: send(pid, {:push, encode(resumed_frame(nil))})

      :manual ->
        :ok

      {:close, code} ->
        send(pid, {:close, code})

      {:invalid, d} ->
        send(pid, {:push, encode(%{"op" => 9, "d" => d})})
    end
  end

  defp react(_state, _pid, _map), do: :ok

  # --- frame builders -----------------------------------------------------

  defp push(state, map) do
    if state.current, do: send(state.current, {:push, encode(map)})
  end

  defp encode(map), do: {:text, JSON.encode!(map)}

  defp hello_frame(interval),
    do: %{"op" => 10, "d" => %{"heartbeat_interval" => interval}}

  defp ready_frame(sid, resume_url) do
    %{
      "op" => 0,
      "t" => "READY",
      "s" => 1,
      "d" => %{"session_id" => sid, "resume_gateway_url" => resume_url}
    }
  end

  defp resumed_frame(seq) do
    base = %{"op" => 0, "t" => "RESUMED", "d" => %{}}
    if is_integer(seq), do: Map.put(base, "s", seq), else: base
  end

  defp dispatch_frame(name, data, seq),
    do: %{"op" => 0, "t" => name, "s" => seq, "d" => data}

  defp ws_url(state), do: "ws://127.0.0.1:#{state.port}"
end