From e9b711d7134337f29f406e66d5e36a342cbf9352 Mon Sep 17 00:00:00 2001
From: Mat Trudel <mat@geeky.net>
Date: Sun, 5 Jan 2025 12:45:05 -0500
Subject: [PATCH] Improve tests (#149)

* Move telemetry tests to use message based helper from Bandit

* Use async tests
---
 lib/thousand_island/listener.ex        |   5 +-
 lib/thousand_island/telemetry.ex       |  10 +-
 test/support/telemetry_collector.ex    |  38 -----
 test/support/telemetry_helpers.ex      |  34 ++++
 test/thousand_island/handler_test.exs  | 206 ++++++++++++-------------
 test/thousand_island/listener_test.exs |  79 ++++------
 test/thousand_island/server_test.exs   | 109 +++++++------
 test/thousand_island/socket_test.exs   |  33 ++--
 8 files changed, 251 insertions(+), 263 deletions(-)
 delete mode 100644 test/support/telemetry_collector.ex
 create mode 100644 test/support/telemetry_helpers.ex

diff --git a/lib/thousand_island/listener.ex b/lib/thousand_island/listener.ex
index d9e18c1..07c50be 100644
--- a/lib/thousand_island/listener.ex
+++ b/lib/thousand_island/listener.ex
@@ -32,14 +32,15 @@ defmodule ThousandIsland.Listener do
            ),
          {:ok, {ip, port}} <-
            server_config.transport_module.sockname(listener_socket) do
-      span_meta = %{
+      span_metadata = %{
+        handler: server_config.handler_module,
         local_address: ip,
         local_port: port,
         transport_module: server_config.transport_module,
         transport_options: server_config.transport_options
       }
 
-      listener_span = ThousandIsland.Telemetry.start_span(:listener, %{}, span_meta)
+      listener_span = ThousandIsland.Telemetry.start_span(:listener, %{}, span_metadata)
 
       {:ok,
        %{listener_socket: listener_socket, local_info: {ip, port}, listener_span: listener_span}}
diff --git a/lib/thousand_island/telemetry.ex b/lib/thousand_island/telemetry.ex
index 187a3e8..6218e7f 100644
--- a/lib/thousand_island/telemetry.ex
+++ b/lib/thousand_island/telemetry.ex
@@ -332,7 +332,9 @@ defmodule ThousandIsland.Telemetry do
   @spec start_child_span(t(), span_name(), measurements(), metadata()) :: t()
   def start_child_span(parent_span, span_name, measurements \\ %{}, metadata \\ %{}) do
     metadata =
-      Map.put(metadata, :parent_telemetry_span_context, parent_span.telemetry_span_context)
+      metadata
+      |> Map.put(:parent_telemetry_span_context, parent_span.telemetry_span_context)
+      |> Map.put(:handler, parent_span.start_metadata.handler)
 
     start_span(span_name, measurements, metadata)
   end
@@ -361,7 +363,11 @@ defmodule ThousandIsland.Telemetry do
   @spec untimed_span_event(t(), event_name() | untimed_event_name(), measurements(), metadata()) ::
           :ok
   def untimed_span_event(span, name, measurements \\ %{}, metadata \\ %{}) do
-    metadata = Map.put(metadata, :telemetry_span_context, span.telemetry_span_context)
+    metadata =
+      metadata
+      |> Map.put(:telemetry_span_context, span.telemetry_span_context)
+      |> Map.put(:handler, span.start_metadata.handler)
+
     event([span.span_name, name], measurements, metadata)
   end
 
diff --git a/test/support/telemetry_collector.ex b/test/support/telemetry_collector.ex
deleted file mode 100644
index 6831eb0..0000000
--- a/test/support/telemetry_collector.ex
+++ /dev/null
@@ -1,38 +0,0 @@
-defmodule ThousandIsland.TelemetryCollector do
-  @moduledoc false
-
-  use GenServer
-
-  def start_link(event_names) do
-    GenServer.start_link(__MODULE__, event_names)
-  end
-
-  def record_event(event, measurements, metadata, pid) do
-    GenServer.cast(pid, {:event, event, measurements, metadata})
-  end
-
-  def get_events(pid) do
-    GenServer.call(pid, :get_events)
-  end
-
-  def init(event_names) do
-    # Use __MODULE__ here to keep telemetry from warning about passing a local capture
-    # https://hexdocs.pm/telemetry/telemetry.html#attach/4
-    :telemetry.attach_many(
-      "#{inspect(self())}.trace",
-      event_names,
-      &__MODULE__.record_event/4,
-      self()
-    )
-
-    {:ok, []}
-  end
-
-  def handle_cast({:event, event, measurements, metadata}, events) do
-    {:noreply, [{event, measurements, metadata} | events]}
-  end
-
-  def handle_call(:get_events, _from, events) do
-    {:reply, Enum.reverse(events), events}
-  end
-end
diff --git a/test/support/telemetry_helpers.ex b/test/support/telemetry_helpers.ex
new file mode 100644
index 0000000..8841cb4
--- /dev/null
+++ b/test/support/telemetry_helpers.ex
@@ -0,0 +1,34 @@
+defmodule TelemetryHelpers do
+  @moduledoc false
+
+  @events [
+    [:thousand_island, :listener, :start],
+    [:thousand_island, :listener, :stop],
+    [:thousand_island, :acceptor, :start],
+    [:thousand_island, :acceptor, :stop],
+    [:thousand_island, :acceptor, :spawn_error],
+    [:thousand_island, :acceptor, :econnaborted],
+    [:thousand_island, :connection, :start],
+    [:thousand_island, :connection, :stop],
+    [:thousand_island, :connection, :ready],
+    [:thousand_island, :connection, :async_recv],
+    [:thousand_island, :connection, :recv],
+    [:thousand_island, :connection, :recv_error],
+    [:thousand_island, :connection, :send],
+    [:thousand_island, :connection, :send_error],
+    [:thousand_island, :connection, :sendfile],
+    [:thousand_island, :connection, :sendfile_error],
+    [:thousand_island, :connection, :socket_shutdown]
+  ]
+
+  def attach_all_events(handler) do
+    ref = make_ref()
+    _ = :telemetry.attach_many(ref, @events, &__MODULE__.handle_event/4, {self(), handler})
+    fn -> :telemetry.detach(ref) end
+  end
+
+  def handle_event(event, measurements, %{handler: handler} = metadata, {pid, handler}),
+    do: send(pid, {:telemetry, event, measurements, metadata})
+
+  def handle_event(_event, _measurements, _metadata, {_pid, _handler}), do: :ok
+end
diff --git a/test/thousand_island/handler_test.exs b/test/thousand_island/handler_test.exs
index c533c8c..c15566f 100644
--- a/test/thousand_island/handler_test.exs
+++ b/test/thousand_island/handler_test.exs
@@ -1,6 +1,5 @@
 defmodule ThousandIsland.HandlerTest do
-  # False due to telemetry raciness
-  use ExUnit.Case, async: false
+  use ExUnit.Case, async: true
 
   use Machete
 
@@ -647,64 +646,63 @@ defmodule ThousandIsland.HandlerTest do
     end
 
     test "it should send `start` telemetry event on startup" do
-      {:ok, collector_pid} =
-        start_supervised(
-          {ThousandIsland.TelemetryCollector, [[:thousand_island, :connection, :start]]}
-        )
+      TelemetryHelpers.attach_all_events(Telemetry.Closer)
 
       {:ok, port} = start_handler(Telemetry.Closer)
 
       {:ok, client} = :gen_tcp.connect(:localhost, port, active: false)
       {:ok, {ip, port}} = :inet.sockname(client)
-      Process.sleep(100)
 
-      assert ThousandIsland.TelemetryCollector.get_events(collector_pid)
-             ~> [
-               {[:thousand_island, :connection, :start], %{monotonic_time: integer()},
-                %{
-                  parent_telemetry_span_context: reference(),
-                  remote_address: ip,
-                  remote_port: port,
-                  telemetry_span_context: reference()
-                }}
-             ]
+      assert_receive {:telemetry, [:thousand_island, :connection, :start], measurements,
+                      metadata},
+                     500
+
+      assert measurements ~> %{monotonic_time: integer()}
+
+      assert metadata
+             ~> %{
+               handler: Telemetry.Closer,
+               parent_telemetry_span_context: reference(),
+               remote_address: ip,
+               remote_port: port,
+               telemetry_span_context: reference()
+             }
     end
 
     test "it should send `ready` telemetry event once socket is ready" do
-      {:ok, collector_pid} =
-        start_supervised(
-          {ThousandIsland.TelemetryCollector, [[:thousand_island, :connection, :ready]]}
-        )
+      TelemetryHelpers.attach_all_events(Telemetry.Closer)
 
       {:ok, port} = start_handler(Telemetry.Closer)
 
       {:ok, _client} = :gen_tcp.connect(:localhost, port, active: false)
-      Process.sleep(100)
 
-      assert ThousandIsland.TelemetryCollector.get_events(collector_pid)
-             ~> [
-               {[:thousand_island, :connection, :ready], %{monotonic_time: integer()},
-                %{telemetry_span_context: reference()}}
-             ]
+      assert_receive {:telemetry, [:thousand_island, :connection, :ready], measurements,
+                      metadata},
+                     500
+
+      assert measurements ~> %{monotonic_time: integer()}
+      assert metadata ~> %{handler: Telemetry.Closer, telemetry_span_context: reference()}
     end
 
     test "it should send `async_recv` telemetry event on async receipt of data" do
-      {:ok, collector_pid} =
-        start_supervised(
-          {ThousandIsland.TelemetryCollector, [[:thousand_island, :connection, :async_recv]]}
-        )
+      TelemetryHelpers.attach_all_events(Telemetry.CloseOnData)
 
       {:ok, port} = start_handler(Telemetry.CloseOnData)
 
       {:ok, client} = :gen_tcp.connect(:localhost, port, active: false)
       :gen_tcp.send(client, "ping")
-      Process.sleep(100)
 
-      assert ThousandIsland.TelemetryCollector.get_events(collector_pid)
-             ~> [
-               {[:thousand_island, :connection, :async_recv], %{data: "ping"},
-                %{telemetry_span_context: reference()}}
-             ]
+      assert_receive {:telemetry, [:thousand_island, :connection, :async_recv], measurements,
+                      metadata},
+                     500
+
+      assert measurements ~> %{data: "ping"}
+
+      assert metadata
+             ~> %{
+               handler: Telemetry.CloseOnData,
+               telemetry_span_context: reference()
+             }
     end
 
     defmodule Telemetry.HelloWorld do
@@ -718,10 +716,7 @@ defmodule ThousandIsland.HandlerTest do
     end
 
     test "it should send `stop` telemetry event on client shutdown" do
-      {:ok, collector_pid} =
-        start_supervised(
-          {ThousandIsland.TelemetryCollector, [[:thousand_island, :connection, :stop]]}
-        )
+      TelemetryHelpers.attach_all_events(Telemetry.HelloWorld)
 
       {:ok, port} = start_handler(Telemetry.HelloWorld)
 
@@ -729,23 +724,20 @@ defmodule ThousandIsland.HandlerTest do
       {:ok, {ip, port}} = :inet.sockname(client)
       {:ok, ~c"HELLO"} = :gen_tcp.recv(client, 5)
       :gen_tcp.close(client)
-      Process.sleep(100)
-
-      # When the remote end closes the socket we do not get octet counts, unfortunately
-      assert ThousandIsland.TelemetryCollector.get_events(collector_pid)
-             ~> [
-               {[:thousand_island, :connection, :stop],
-                %{
-                  duration: integer(),
-                  monotonic_time: integer()
-                },
-                %{
-                  parent_telemetry_span_context: reference(),
-                  remote_address: ip,
-                  remote_port: port,
-                  telemetry_span_context: reference()
-                }}
-             ]
+
+      assert_receive {:telemetry, [:thousand_island, :connection, :stop], measurements, metadata},
+                     500
+
+      assert measurements ~> %{monotonic_time: integer(), duration: integer()}
+
+      assert metadata
+             ~> %{
+               handler: Telemetry.HelloWorld,
+               parent_telemetry_span_context: reference(),
+               remote_address: ip,
+               remote_port: port,
+               telemetry_span_context: reference()
+             }
     end
 
     defmodule Telemetry.Closer do
@@ -759,10 +751,7 @@ defmodule ThousandIsland.HandlerTest do
     end
 
     test "it should send `stop` telemetry event on shutdown" do
-      {:ok, collector_pid} =
-        start_supervised(
-          {ThousandIsland.TelemetryCollector, [[:thousand_island, :connection, :stop]]}
-        )
+      TelemetryHelpers.attach_all_events(Telemetry.Closer)
 
       {:ok, port} = start_handler(Telemetry.Closer)
 
@@ -770,26 +759,28 @@ defmodule ThousandIsland.HandlerTest do
       {:ok, {ip, port}} = :inet.sockname(client)
 
       :gen_tcp.connect(:localhost, port, active: false)
-      Process.sleep(100)
-
-      assert ThousandIsland.TelemetryCollector.get_events(collector_pid)
-             ~> [
-               {[:thousand_island, :connection, :stop],
-                %{
-                  duration: integer(),
-                  monotonic_time: integer(),
-                  recv_cnt: 0,
-                  recv_oct: 0,
-                  send_cnt: 1,
-                  send_oct: 5
-                },
-                %{
-                  parent_telemetry_span_context: reference(),
-                  remote_address: ip,
-                  remote_port: port,
-                  telemetry_span_context: reference()
-                }}
-             ]
+
+      assert_receive {:telemetry, [:thousand_island, :connection, :stop], measurements, metadata},
+                     500
+
+      assert measurements
+             ~> %{
+               monotonic_time: integer(),
+               duration: integer(),
+               recv_cnt: 0,
+               recv_oct: 0,
+               send_cnt: 1,
+               send_oct: 5
+             }
+
+      assert metadata
+             ~> %{
+               handler: Telemetry.Closer,
+               parent_telemetry_span_context: reference(),
+               remote_address: ip,
+               remote_port: port,
+               telemetry_span_context: reference()
+             }
     end
 
     defmodule Telemetry.Error do
@@ -803,41 +794,40 @@ defmodule ThousandIsland.HandlerTest do
 
     @tag capture_log: true
     test "it should send `stop` telemetry event with payload on error" do
-      {:ok, collector_pid} =
-        start_supervised(
-          {ThousandIsland.TelemetryCollector, [[:thousand_island, :connection, :stop]]}
-        )
+      TelemetryHelpers.attach_all_events(Telemetry.Error)
 
       {:ok, port} = start_handler(Telemetry.Error)
 
       {:ok, client} = :gen_tcp.connect(:localhost, port, active: false)
       {:ok, {ip, port}} = :inet.sockname(client)
-      Process.sleep(100)
-
-      assert ThousandIsland.TelemetryCollector.get_events(collector_pid)
-             ~> [
-               {[:thousand_island, :connection, :stop],
-                %{
-                  duration: integer(),
-                  monotonic_time: integer(),
-                  recv_cnt: 0,
-                  recv_oct: 0,
-                  send_cnt: 0,
-                  send_oct: 0
-                },
-                %{
-                  error: :nope,
-                  parent_telemetry_span_context: reference(),
-                  remote_address: ip,
-                  remote_port: port,
-                  telemetry_span_context: reference()
-                }}
-             ]
+
+      assert_receive {:telemetry, [:thousand_island, :connection, :stop], measurements, metadata},
+                     500
+
+      assert measurements
+             ~> %{
+               monotonic_time: integer(),
+               duration: integer(),
+               recv_cnt: 0,
+               recv_oct: 0,
+               send_cnt: 0,
+               send_oct: 0
+             }
+
+      assert metadata
+             ~> %{
+               handler: Telemetry.Error,
+               error: :nope,
+               parent_telemetry_span_context: reference(),
+               remote_address: ip,
+               remote_port: port,
+               telemetry_span_context: reference()
+             }
     end
   end
 
   defp start_handler(handler, server_args \\ []) do
-    resolved_args = [port: 0, handler_module: handler] ++ server_args
+    resolved_args = [port: 0, handler_module: handler, num_acceptors: 1] ++ server_args
     {:ok, server_pid} = start_supervised({ThousandIsland, resolved_args})
     {:ok, {_ip, port}} = ThousandIsland.listener_info(server_pid)
     {:ok, port}
diff --git a/test/thousand_island/listener_test.exs b/test/thousand_island/listener_test.exs
index c11afa0..8d7e054 100644
--- a/test/thousand_island/listener_test.exs
+++ b/test/thousand_island/listener_test.exs
@@ -1,10 +1,11 @@
 defmodule ThousandIsland.ListenerTest do
-  use ExUnit.Case, async: false
+  use ExUnit.Case, async: true
   use Machete
 
   alias ThousandIsland.{Listener, ServerConfig}
 
-  @server_config %ServerConfig{port: 4004}
+  # We don't actually implement handler, but we specify it so that our telemetry helpers will work
+  @server_config %ServerConfig{port: 4004, handler_module: __MODULE__}
 
   defmodule TestTransport do
     # This module does not implement all of the callbacks
@@ -90,29 +91,24 @@ defmodule ThousandIsland.ListenerTest do
     end
 
     test "emits expected telemetry event" do
-      {:ok, collector_pid} =
-        start_supervised(
-          {ThousandIsland.TelemetryCollector,
-           [
-             [:thousand_island, :listener, :start]
-           ]}
-        )
+      TelemetryHelpers.attach_all_events(__MODULE__)
 
-      {:ok, %{listener_socket: socket}} =
-        Listener.init(@server_config)
+      {:ok, %{listener_socket: socket}} = Listener.init(@server_config)
+
+      assert_receive {:telemetry, [:thousand_island, :listener, :start], measurements, metadata},
+                     500
+
+      assert measurements ~> %{monotonic_time: integer()}
 
-      # We expect a monotonic start time as a measurement in the event.
-      assert ThousandIsland.TelemetryCollector.get_events(collector_pid)
-             ~> [
-               {[:thousand_island, :listener, :start], %{monotonic_time: integer()},
-                %{
-                  telemetry_span_context: reference(),
-                  local_address: {0, 0, 0, 0},
-                  local_port: @server_config.port,
-                  transport_module: ThousandIsland.Transports.TCP,
-                  transport_options: []
-                }}
-             ]
+      assert metadata
+             ~> %{
+               handler: __MODULE__,
+               telemetry_span_context: reference(),
+               local_address: {0, 0, 0, 0},
+               local_port: @server_config.port,
+               transport_module: ThousandIsland.Transports.TCP,
+               transport_options: []
+             }
 
       # Close the socket to cleanup.
       :gen_tcp.close(socket)
@@ -145,35 +141,26 @@ defmodule ThousandIsland.ListenerTest do
 
   describe "terminate/2" do
     test "emits telemetry event with expected timings" do
-      {:ok, %{listener_span: span, listener_socket: socket}} =
-        Listener.init(@server_config)
+      {:ok, %{listener_span: span, listener_socket: socket}} = Listener.init(@server_config)
 
-      {:ok, collector_pid} =
-        start_supervised(
-          {ThousandIsland.TelemetryCollector,
-           [
-             [:thousand_island, :listener, :stop]
-           ]}
-        )
+      TelemetryHelpers.attach_all_events(__MODULE__)
 
       Listener.terminate(:normal, %{listener_span: span})
 
-      # We expect a monotonic stop time and a duration
-      # as measurements in the event.
-      assert [
-               {[:thousand_island, :listener, :stop],
-                %{monotonic_time: stop_monotonic_time, duration: duration}, stop_metadata}
-             ] = ThousandIsland.TelemetryCollector.get_events(collector_pid)
-
-      assert is_integer(stop_monotonic_time)
+      assert_receive {:telemetry, [:thousand_island, :listener, :stop], measurements, metadata},
+                     500
 
-      # We expect the duration to be the monotonic stop
-      # time minus the monotonic start time.
-      assert stop_monotonic_time >= span.start_time
-      assert duration == stop_monotonic_time - span.start_time
+      assert measurements ~> %{monotonic_time: integer(), duration: integer()}
 
-      # The start and stop metadata should be equal.
-      assert stop_metadata == span.start_metadata
+      assert metadata
+             ~> %{
+               handler: __MODULE__,
+               telemetry_span_context: reference(),
+               local_address: {0, 0, 0, 0},
+               local_port: @server_config.port,
+               transport_module: ThousandIsland.Transports.TCP,
+               transport_options: []
+             }
 
       # Close the socket to cleanup.
       :gen_tcp.close(socket)
diff --git a/test/thousand_island/server_test.exs b/test/thousand_island/server_test.exs
index 473ec6e..1e4145a 100644
--- a/test/thousand_island/server_test.exs
+++ b/test/thousand_island/server_test.exs
@@ -1,6 +1,5 @@
 defmodule ThousandIsland.ServerTest do
-  # False due to telemetry raciness
-  use ExUnit.Case, async: false
+  use ExUnit.Case, async: true
 
   use Machete
 
@@ -162,10 +161,7 @@ defmodule ThousandIsland.ServerTest do
     end
 
     test "should emit telemetry events as expected" do
-      {:ok, collector_pid} =
-        start_supervised(
-          {ThousandIsland.TelemetryCollector, [[:thousand_island, :acceptor, :spawn_error]]}
-        )
+      TelemetryHelpers.attach_all_events(LongEcho)
 
       {:ok, _, port} =
         start_handler(LongEcho,
@@ -180,14 +176,12 @@ defmodule ThousandIsland.ServerTest do
       :ok = :gen_tcp.send(client, "HELLO")
       :ok = :gen_tcp.send(other_client, "BONJOUR")
 
-      # Give things enough time for the second connection to time out
-      Process.sleep(700)
+      assert_receive {:telemetry, [:thousand_island, :acceptor, :spawn_error], measurements,
+                      metadata},
+                     1000
 
-      assert ThousandIsland.TelemetryCollector.get_events(collector_pid)
-             ~> [
-               {[:thousand_island, :acceptor, :spawn_error], %{monotonic_time: integer()},
-                %{telemetry_span_context: reference()}}
-             ]
+      assert measurements ~> %{monotonic_time: integer()}
+      assert metadata ~> %{handler: LongEcho, telemetry_span_context: reference()}
     end
   end
 
@@ -332,46 +326,65 @@ defmodule ThousandIsland.ServerTest do
     end
 
     test "it should emit telemetry events as expected" do
-      {:ok, collector_pid} =
-        start_supervised(
-          {ThousandIsland.TelemetryCollector,
-           [
-             [:thousand_island, :listener, :start],
-             [:thousand_island, :listener, :stop],
-             [:thousand_island, :acceptor, :start],
-             [:thousand_island, :acceptor, :stop]
-           ]}
-        )
+      TelemetryHelpers.attach_all_events(Echo)
 
       {:ok, server_pid, _} = start_handler(Echo, num_acceptors: 1)
 
+      assert_receive {:telemetry, [:thousand_island, :listener, :start], measurements, metadata},
+                     500
+
+      assert measurements ~> %{monotonic_time: integer()}
+
+      assert metadata
+             ~> %{
+               handler: Echo,
+               telemetry_span_context: reference(),
+               local_address: {0, 0, 0, 0},
+               local_port: integer(),
+               transport_module: ThousandIsland.Transports.TCP,
+               transport_options: []
+             }
+
+      assert_receive {:telemetry, [:thousand_island, :acceptor, :start], measurements, metadata},
+                     500
+
+      assert measurements ~> %{monotonic_time: integer()}
+
+      assert metadata
+             ~> %{
+               handler: Echo,
+               telemetry_span_context: reference(),
+               parent_telemetry_span_context: reference()
+             }
+
       ThousandIsland.stop(server_pid)
 
-      assert ThousandIsland.TelemetryCollector.get_events(collector_pid)
-             ~> [
-               {[:thousand_island, :listener, :start], %{monotonic_time: integer()},
-                %{
-                  telemetry_span_context: reference(),
-                  local_address: {0, 0, 0, 0},
-                  local_port: integer(),
-                  transport_module: ThousandIsland.Transports.TCP,
-                  transport_options: []
-                }},
-               {[:thousand_island, :acceptor, :start], %{monotonic_time: integer()},
-                %{telemetry_span_context: reference(), parent_telemetry_span_context: reference()}},
-               {[:thousand_island, :listener, :stop],
-                %{duration: integer(), monotonic_time: integer()},
-                %{
-                  telemetry_span_context: reference(),
-                  local_address: {0, 0, 0, 0},
-                  local_port: integer(),
-                  transport_module: ThousandIsland.Transports.TCP,
-                  transport_options: []
-                }},
-               {[:thousand_island, :acceptor, :stop],
-                %{connections: 0, duration: integer(), monotonic_time: integer()},
-                %{telemetry_span_context: reference(), parent_telemetry_span_context: reference()}}
-             ]
+      assert_receive {:telemetry, [:thousand_island, :acceptor, :stop], measurements, metadata},
+                     500
+
+      assert measurements ~> %{monotonic_time: integer(), duration: integer(), connections: 0}
+
+      assert metadata
+             ~> %{
+               handler: Echo,
+               telemetry_span_context: reference(),
+               parent_telemetry_span_context: reference()
+             }
+
+      assert_receive {:telemetry, [:thousand_island, :listener, :stop], measurements, metadata},
+                     500
+
+      assert measurements ~> %{monotonic_time: integer(), duration: integer()}
+
+      assert metadata
+             ~> %{
+               handler: Echo,
+               telemetry_span_context: reference(),
+               local_address: {0, 0, 0, 0},
+               local_port: integer(),
+               transport_module: ThousandIsland.Transports.TCP,
+               transport_options: []
+             }
     end
   end
 
diff --git a/test/thousand_island/socket_test.exs b/test/thousand_island/socket_test.exs
index 632bd38..a94fb2a 100644
--- a/test/thousand_island/socket_test.exs
+++ b/test/thousand_island/socket_test.exs
@@ -1,6 +1,5 @@
 defmodule ThousandIsland.SocketTest do
-  # False due to telemetry raciness
-  use ExUnit.Case, async: false
+  use ExUnit.Case, async: true
 
   use Machete
 
@@ -99,14 +98,7 @@ defmodule ThousandIsland.SocketTest do
       end
 
       test "it should emit telemetry events as expected", context do
-        {:ok, collector_pid} =
-          start_supervised(
-            {ThousandIsland.TelemetryCollector,
-             [
-               [:thousand_island, :connection, :recv],
-               [:thousand_island, :connection, :send]
-             ]}
-          )
+        TelemetryHelpers.attach_all_events(Echo)
 
         {:ok, port} = start_handler(Echo, context.server_opts)
         {:ok, client} = context.client_mod.connect(~c"localhost", port, context.client_opts)
@@ -115,16 +107,19 @@ defmodule ThousandIsland.SocketTest do
         {:ok, ~c"HELLO"} = context.client_mod.recv(client, 0)
         context.client_mod.close(client)
 
-        # Give the server process a chance to shut down
-        Process.sleep(100)
+        assert_receive {:telemetry, [:thousand_island, :connection, :recv], measurements,
+                        metadata},
+                       500
 
-        assert ThousandIsland.TelemetryCollector.get_events(collector_pid)
-               ~> [
-                 {[:thousand_island, :connection, :recv], %{data: "HELLO"},
-                  %{telemetry_span_context: reference()}},
-                 {[:thousand_island, :connection, :send], %{data: "HELLO"},
-                  %{telemetry_span_context: reference()}}
-               ]
+        assert measurements ~> %{data: "HELLO"}
+        assert metadata ~> %{handler: Echo, telemetry_span_context: reference()}
+
+        assert_receive {:telemetry, [:thousand_island, :connection, :send], measurements,
+                        metadata},
+                       500
+
+        assert measurements ~> %{data: "HELLO"}
+        assert metadata ~> %{handler: Echo, telemetry_span_context: reference()}
       end
     end
   end)