Define TestSocket for test on Phoenix 1.4

I used to define TestSocket on tests to test some channels. (There are several sockets in my project)

defmodule MyWeb.SomeTest do
  use ExchangeWeb.ChannelCase

  defmodule TestSocket do
    use ExchangeWeb, :socket

    channel("/test_channel*", TestChannel)

    transport(:websocket, Phoenix.Transports.WebSocket)

    def connect(_params, socket) do
        {:ok, socket}
    end

    def id(_socket) do
      "test_socket:1"
    end
  end

  test(...)
end

And I need mock socket much more as Phoenix.ChannelTest.socket/0, 2 are deprecated.
But the way to define socket is moved from socket to endpoint.

defmodule MyWeb.Endpoint do
  socket "/socket", MyWeb.UserSocket,
    websocket: true,
    longpoll: false
  ...
end

I don’t know how to define test socket on Phoenix 1.4

If there is no other way

defmodule MyWeb.Endpoint do
  if Mix.env == :test do
    socket "/test_socket", MyWeb.TestSocket, websocket: true, longpoll: false
  end
end

might work.