       import Wallaby.Feature
 
       setup context do
-        metadata = configure_ecto(unquote(@includes_ecto), context[:async])
+        metadata = unquote(__MODULE__).Utils.maybe_checkout_repos(context[:async])
 
         start_session_opts =
           [metadata: metadata]
-          |> put_create_session_fn(context[:create_session_fn])
-
-        sessions =
-          get_in(context, [:registered, :sessions])
-          |> sessions_iterable()
-          |> Enum.map(&start_session(&1, start_session_opts))
-
-        sessions |> build_setup_return()
-      end
-    end
-  end
-
-  @doc false
-  def build_setup_return([session] = sessions) when length(sessions) == 1 do
-    [session: session]
-  end
-
-  def build_setup_return(sessions) do
-    [sessions: sessions]
-  end
-
-  @doc false
-  def sessions_iterable(nil), do: 1..1
-  def sessions_iterable(count) when is_number(count), do: 1..count
-  def sessions_iterable(capabilities) when is_list(capabilities), do: capabilities
-
-  @doc false
-  def start_session(more_opts, start_session_opts) when is_list(more_opts) do
-    {:ok, session} =
-      start_session_opts
-      |> Keyword.merge(more_opts)
-      |> Wallaby.start_session()
-
-    session
-  end
-
-  def start_session(num, start_session_opts) when is_number(num) do
-    {:ok, session} = Wallaby.start_session(start_session_opts)
-
-    session
-  end
-
-  @doc false
-  def put_create_session_fn(opts, nil), do: opts
-  def put_create_session_fn(opts, func), do: Keyword.put(opts, :create_session_fn, func)
-
-  if @includes_ecto do
-    @doc false
-    def otp_app(), do: Application.get_env(:wallaby, :otp_app)
-    @doc false
-    def ecto_repos(nil), do: []
-    def ecto_repos(otp_app), do: Application.get_env(otp_app, :ecto_repos, [])
-
-    @doc false
-    def checkout_ecto_repos(repo, async) do
-      :ok = Ecto.Adapters.SQL.Sandbox.checkout(repo)
-
-      unless async, do: Ecto.Adapters.SQL.Sandbox.mode(repo, {:shared, self()})
-
-      repo
-    end
-
-    @doc false
-    def metadata_for_ecto_repos([]), do: Map.new()
-
-    def metadata_for_ecto_repos(repos) do
-      Phoenix.Ecto.SQL.Sandbox.metadata_for(repos, self())
-    end
-  end
-
-  @doc false
-  defmacro configure_ecto(includes_ecto?, async?) do
-    if includes_ecto? do
-      quote do
-        otp_app()
-        |> ecto_repos()
-        |> Enum.map(&checkout_ecto_repos(&1, unquote(async?)))
-        |> metadata_for_ecto_repos()
-      end
-    else
-      quote do
-        ""
+          |> unquote(__MODULE__).Utils.put_create_session_fn(context[:create_session_fn])
+
+        get_in(context, [:registered, :sessions])
+        |> unquote(__MODULE__).Utils.sessions_iterable()
+        |> Enum.map(fn
+          opts when is_list(opts) ->
+            unquote(__MODULE__).Utils.start_session(opts, start_session_opts)
+
+          i when is_number(i) ->
+            unquote(__MODULE__).Utils.start_session([], start_session_opts)
+        end)
+        |> unquote(__MODULE__).Utils.build_setup_return()
       end
     end
   end

        rescue
          e ->
            if Wallaby.screenshot_on_failure?() do
              unquote(__MODULE__).take_screenshots_for_sessions(self(), unquote(message))
            end

            reraise(e, __STACKTRACE__)

    end
  end

  defmodule Utils do
    @includes_ecto Code.ensure_loaded?(Ecto.Adapters.SQL.Sandbox) &&
                     Code.ensure_loaded?(Phoenix.Ecto.SQL.Sandbox)
    @moduledoc false

    def build_setup_return([session]) do
      [session: session]
    end

    def build_setup_return(sessions) do
      [sessions: sessions]
    end

    def sessions_iterable(nil), do: 1..1
    def sessions_iterable(count) when is_number(count), do: 1..count
    def sessions_iterable(capabilities) when is_list(capabilities), do: capabilities

    def start_session(more_opts, start_session_opts) when is_list(more_opts) do
      {:ok, session} =
        start_session_opts
        |> Keyword.merge(more_opts)
        |> Wallaby.start_session()

      session
    end

    def put_create_session_fn(opts, nil), do: opts
    def put_create_session_fn(opts, func), do: Keyword.put(opts, :create_session_fn, func)

    if @includes_ecto do
      def maybe_checkout_repos(async?) do
        otp_app()
        |> ecto_repos()
        |> Enum.map(&checkout_ecto_repos(&1, async?))
        |> metadata_for_ecto_repos()
      end

      defp otp_app(), do: Application.get_env(:wallaby, :otp_app)

      defp ecto_repos(nil), do: []
      defp ecto_repos(otp_app), do: Application.get_env(otp_app, :ecto_repos, [])

      defp checkout_ecto_repos(repo, async) do
        :ok = Ecto.Adapters.SQL.Sandbox.checkout(repo)

        unless async, do: Ecto.Adapters.SQL.Sandbox.mode(repo, {:shared, self()})

        repo
      end

      defp metadata_for_ecto_repos([]), do: Map.new()

      defp metadata_for_ecto_repos(repos) do
        Phoenix.Ecto.SQL.Sandbox.metadata_for(repos, self())
      end
    else
      def maybe_checkout_repos(_) do
        ""
      end
    end

    def take_screenshots_for_sessions(pid, test_name) do
      time = :erlang.system_time(:second) |> to_string()
      test_name = String.replace(test_name, " ", "_")

      screenshot_paths =
        Wallaby.SessionStore.list_sessions_for(pid)
        |> Enum.with_index(1)
        |> Enum.flat_map(fn {s, i} ->
          filename = time <> "_" <> test_name <> "(#{i})"

          Wallaby.Browser.take_screenshot(s, name: filename).screenshots
        end)
        |> Enum.map_join("\n- ", &Wallaby.Browser.build_file_url/1)

      IO.write("\n- #{screenshot_paths}")
    end
  end
end
