Add TestServer.TCP

Author: richard-ash

README.md

@@ -13,13 +13,15 @@ Features:
 - HTTP/2
 - WebSocket
 - SSH
+- TCP
 - Built-in TLS with self-signed certificates
 - Plug route matching
 
 ## Protocols
 
 - [`TestServer.HTTP`](lib/test_server/http/README.md) - HTTP/1, HTTP/2, and WebSocket.
 - [`TestServer.SSH`](lib/test_server/ssh/README.md) - SSH exec and shell.
+- [`TestServer.TCP`](lib/test_server/tcp/README.md) - Raw TCP.
 
 <!-- MDOC !-->
 

lib/test_server/tcp.ex

@@ -0,0 +1,231 @@
+defmodule TestServer.TCP do
+  @external_resource "lib/test_server/tcp/README.md"
+  @moduledoc "lib/test_server/tcp/README.md"
+             |> File.read!()
+             |> String.split("<!-- MDOC !-->")
+             |> Enum.fetch!(1)
+
+  alias TestServer.TCP.{Instance, Server}
+
+  @type connection :: {pid(), connection_ref()}
+  @type connection_ref :: reference()
+  @type data :: binary()
+  @type state :: term()
+  @type handler_fun :: (data(), state() ->
+                          {:reply, iodata(), state()}
+                          | {:ok, state()}
+                          | {:close, state()})
+  @type raw_handler_fun :: (data(), port(), state() ->
+                              {:ok, state()}
+                              | {:close, state()})
+  @type match_fun :: (data(), state() -> boolean())
+
+  @doc """
+  Start a test server TCP instance.
+
+  The instance will be terminated when the test case finishes.
+
+  ## Options
+
+    * `:port`           - integer of port number, defaults to random port that
+      can be opened;
+    * `:ipfamily`       - The IP address type to use, either `:inet` or
+      `:inet6`. Defaults to `:inet`;
+    * `:listen_options` - options passed to `:gen_tcp.listen/2`. Defaults to
+      `[:binary, active: false, reuseaddr: true]`. `active: false` is always
+      used by the server;
+    * `:recv_timeout`   - timeout passed to `:gen_tcp.recv/3`. Defaults to
+      `5_000`.
+
+  ## Examples
+
+      {:ok, _instance} = TestServer.TCP.start(
+        listen_options: [:binary, packet: :line]
+      )
+
+      {:ok, connection} = TestServer.TCP.connect()
+
+      :ok =
+        TestServer.TCP.handle(connection,
+          match: fn data, _state -> data == "PING\\n" end,
+          to: fn _data, state -> {:reply, "PONG\\n", state} end
+        )
+
+      {:ok, socket} =
+        :gen_tcp.connect(~c"localhost", elem(TestServer.TCP.address(), 1), [
+          :binary,
+          active: false,
+          packet: :line
+        ])
+
+      :ok = :gen_tcp.send(socket, "PING\\n")
+      assert {:ok, "PONG\\n"} = :gen_tcp.recv(socket, 0)
+  """
+  @spec start(keyword()) :: {:ok, pid()}
+  def start(options \\ []) do
+    TestServer.start_instance(__MODULE__, options, &verify!/1)
+  end
+
+  defp verify!(instance) do
+    verify_handlers!(instance)
+    verify_connections!(instance)
+  end
+
+  defp verify_handlers!(instance) do
+    instance
+    |> Instance.handlers()
+    |> Enum.reject(& &1.suspended)
+    |> case do
+      [] ->
+        :ok
+
+      active_handlers ->
+        raise """
+        #{TestServer.format_instance(__MODULE__, instance)} did not receive data for these handlers before the test ended:
+
+        #{Instance.format_handlers(active_handlers)}
+        """
+    end
+  end
+
+  defp verify_connections!(instance) do
+    instance
+    |> Instance.connections()
+    |> Enum.filter(&is_nil(&1.pid))
+    |> case do
+      [] ->
+        :ok
+
+      unused_connections ->
+        raise """
+        #{TestServer.format_instance(__MODULE__, instance)} has connections that were not used:
+
+        #{Instance.format_connections(unused_connections)}
+        """
+    end
+  end
+
+  @doc """
+  Shuts down the current test server TCP instance.
+  """
+  @spec stop() :: :ok | {:error, term()}
+  def stop, do: stop(TestServer.fetch_instance!(__MODULE__))
+
+  @doc """
+  Shuts down a test server TCP instance.
+  """
+  @spec stop(pid()) :: :ok | {:error, term()}
+  def stop(instance) do
+    TestServer.ensure_instance_alive!(__MODULE__, instance)
+
+    Server.stop(Instance.get_options(instance))
+
+    TestServer.stop_instance(__MODULE__, instance)
+  end
+
+  @spec address() :: {binary(), non_neg_integer()}
+  def address, do: address([])
+
+  @doc """
+  Returns the address for current test server.
+
+  ## Options
+
+    * `:host` - binary host value, it'll be added to inet for IP `127.0.0.1`
+      and `::1`, defaults to `"localhost"`;
+  """
+  @spec address(keyword() | pid()) :: {binary(), non_neg_integer()}
+  def address(options) when is_list(options),
+    do: address(TestServer.fetch_instance!(__MODULE__), options)
+
+  def address(instance) when is_pid(instance), do: address(instance, [])
+
+  @doc """
+  Returns the address for a test server instance.
+
+  See `address/1` for options.
+  """
+  @spec address(pid(), keyword()) :: {binary(), non_neg_integer()}
+  def address(instance, options) when is_pid(instance) and is_list(options) do
+    TestServer.ensure_instance_alive!(__MODULE__, instance)
+
+    host = TestServer.get_host(options)
+    port = instance |> Instance.get_options() |> Keyword.fetch!(:port)
+
+    {host, port}
+  end
+
+  @spec connect() :: {:ok, connection()}
+  def connect, do: connect([])
+
+  @doc """
+  Adds a connection expectation to the current test server.
+
+  ## Options
+
+    * `:init_state` - initial state for handlers on the accepted TCP connection.
+
+  ## Examples
+
+      {:ok, connection} = TestServer.TCP.connect()
+      :ok = TestServer.TCP.handle(connection)
+  """
+  @spec connect(keyword()) :: {:ok, connection()}
+  def connect(options) when is_list(options) do
+    {:ok, instance} = TestServer.autostart_instance(__MODULE__)
+
+    connect(instance, options)
+  end
+
+  @doc """
+  Adds a connection expectation to a test server instance.
+
+  See `connect/1` for options.
+  """
+  @spec connect(pid(), keyword()) :: {:ok, connection()}
+  def connect(instance, options) do
+    TestServer.ensure_instance_alive!(__MODULE__, instance)
+
+    [_first_module_entry | stacktrace] = TestServer.get_pruned_stacktrace(__MODULE__)
+
+    options = Keyword.put_new(options, :init_state, %{})
+
+    {:ok, connection} = Instance.register(instance, {:connection, {options, stacktrace}})
+
+    {:ok, {instance, connection.ref}}
+  end
+
+  @spec handle(connection()) :: :ok
+  def handle(connection), do: handle(connection, [])
+
+  @doc """
+  Adds a data handler to a test server TCP connection.
+
+  Handlers are matched FIFO (first in, first out). Any data not matched by a
+  handler, or any handlers not consumed by data, will raise an error in the test
+  case.
+
+  The `:to` callback can be either a two-arity `t:handler_fun/0` or a
+  three-arity `t:raw_handler_fun/0`. A two-arity handler uses the default TCP
+  handling for replies and closes. A three-arity handler receives the accepted
+  socket and gives you direct control over socket responses.
+
+  ## Options
+
+    * `:match` - a `t:match_fun/0` function that returns a boolean. Defaults to
+      matching anything;
+    * `:to`    - a `t:handler_fun/0` or `t:raw_handler_fun/0` function called
+      when the handler matches. Defaults to echoing the received data.
+  """
+  @spec handle(connection(), keyword()) :: :ok
+  def handle({instance, connection_ref} = _connection, options) do
+    TestServer.ensure_instance_alive!(__MODULE__, instance)
+
+    [_first_module_entry | stacktrace] = TestServer.get_pruned_stacktrace(__MODULE__)
+
+    {:ok, _handler} =
+      Instance.register(instance, {:handle, {connection_ref, options, stacktrace}})
+
+    :ok
+  end
+end

lib/test_server/tcp/README.md

@@ -0,0 +1,104 @@
+# TCP
+
+<!-- MDOC !-->
+
+Mock raw TCP endpoints with connection and data expectations.
+
+## Usage
+
+Add connection expectations with `TestServer.TCP.connect/1`, then add FIFO
+data handlers with `TestServer.TCP.handle/2`:
+
+```elixir
+test "TCP client" do
+  {:ok, connection} = TestServer.TCP.connect()
+
+  :ok =
+    TestServer.TCP.handle(
+      connection,
+      match: fn data, _state -> data == "PING\n" end,
+      to: fn _data, state -> {:reply, "PONG\n", state} end
+    )
+
+  :ok = TestServer.TCP.handle(connection)
+
+  {:ok, socket} =
+    :gen_tcp.connect(~c"localhost", elem(TestServer.TCP.address(), 1), [
+      :binary,
+      active: false,
+      packet: :line
+    ])
+
+  :ok = :gen_tcp.send(socket, "PING\n")
+  assert {:ok, "PONG\n"} = :gen_tcp.recv(socket, 0)
+
+  :ok = :gen_tcp.send(socket, "echo\n")
+  assert {:ok, "echo\n"} = :gen_tcp.recv(socket, 0)
+end
+```
+
+The server autostarts when `connect/1` is called. Start it explicitly when you
+need custom socket options:
+
+```elixir
+TestServer.TCP.start(listen_options: [:binary, packet: :line])
+```
+
+TCP is a stream protocol, so data delivered to handlers follows the configured
+`:gen_tcp` packet options. Use `:listen_options` to set framing such as
+`packet: :line`, `packet: 4`, or raw stream mode.
+
+### Handlers
+
+By default, `TestServer.TCP.handle/2` echoes received data:
+
+```elixir
+{:ok, connection} = TestServer.TCP.connect()
+:ok = TestServer.TCP.handle(connection)
+```
+
+Use `:match` to select data and `:to` to customize the response:
+
+```elixir
+TestServer.TCP.handle(connection,
+  match: fn data, _state -> data == "HELLO" end,
+  to: fn _data, state -> {:reply, "READY", state} end
+)
+```
+
+The two-arity `:to` callback can return:
+
+```elixir
+{:reply, data, state}
+{:ok, state}
+{:close, state}
+```
+
+For lower-level socket control, use a three-arity callback:
+
+```elixir
+TestServer.TCP.handle(connection,
+  to: fn _data, socket, state ->
+    :ok = :gen_tcp.send(socket, "first")
+    :ok = :gen_tcp.send(socket, "second")
+
+    {:ok, state}
+  end
+)
+```
+
+### IPv6
+
+Use the `:ipfamily` option to test with IPv6:
+
+```elixir
+{:ok, _instance} = TestServer.TCP.start(ipfamily: :inet6)
+{:ok, _connection} = TestServer.TCP.connect()
+
+assert {"localhost", port} = TestServer.TCP.address()
+
+{:ok, socket} =
+  :gen_tcp.connect(~c"localhost", port, [:binary, active: false, :inet6])
+```
+
+<!-- MDOC !-->