   end
 
   test "phoenix channel returns heartbeat message when received" do
-    socket = %Socket{pid: self, router: Router9, channel: "phoenix"}
-    sockets = HashDict.put(HashDict.new, {"phoenix", "conn"}, socket)
+    sockets = HashDict.put(HashDict.new, {"phoenix", "conn"}, new_socket)
     message = %Message{channel: "phoenix", topic: "conn", event: "heartbeat", message: %{}}
 
-    assert match?({:ok, _sockets}, Transport.dispatch(message, sockets, self, Router9))
+    assert match?({:ok, _sockets}, Transport.dispatch(message, sockets, self, Router))
     assert_received %Message{channel: "phoenix", topic: "conn", event: "heartbeat", message: %{}}
   end
 
   test "socket state can change when receiving regular process messages" do
-    socket = %Socket{pid: self, router: Router, channel: "chan6", topic: "topic"}
-    sockets = HashDict.put(HashDict.new, {"chan6", "topic"}, socket)
-    message = %Message{channel: "chan6", topic: "topic", event: "join", message: %{}}
-
-    PubSub.create("chan6:topic")
-    {:ok, sockets} = Transport.dispatch(message, sockets, self, Router)
-    {:ok, sockets} = Transport.dispatch_info(sockets, :stuff)
-    socket = HashDict.get(sockets, {"chan6", "topic"})
+    socket = new_socket
+    sockets = HashDict.put(HashDict.new, {"chan1", "topic"}, socket)
+    message = join_message({:ok, socket})
+
+    PubSub.create("chan1:topic")
+    {:ok, sockets} = Transport.dispatch(message, sockets, self, Router)
+    {:ok, sockets} = Transport.dispatch_info(sockets, Socket.assign(socket, :foo, :bar))
+    socket = HashDict.get(sockets, {"chan1", "topic"})
 
     assert socket.assigns[:foo] == :bar
   end
 
   test "Socket state can be put and retrieved" do
-    socket = %Socket{pid: self, router: Router, channel: "chan66"}
-    socket = ChangeSocketStateChannel.event(socket, "put", %{val: 123})
-    _socket = ChangeSocketStateChannel.event(socket, "get", %{"key" => :val})
+    socket = MyChannel.event(new_socket, "put", %{val: 123})
+    _socket = MyChannel.event(socket, "get", %{"key" => :val})
     assert_received 123
   end
 end

  end

  test "unsuccessful join denies socket access to channel/topic" do
    message = %Message{channel: "chan2", topic: "topic", event: "join", message: %{}}

    PubSub.create("chan2:topic")
    assert PubSub.subscribers("chan2:topic") == []
    {:error, sockets, :unauthenticated} = Transport.dispatch(message, HashDict.new, self, Router)
    refute HashDict.get(sockets, {"chan2", "topic"})
    refute PubSub.subscribers("chan2:topic") == [self]
  end

  test "#leave is called when the socket conn closes, and is unsubscribed" do
    socket = %Socket{pid: self, router: Router, channel: "chan1"}
    sockets = HashDict.put(HashDict.new, {"chan1", "topic"}, socket)
    message = %Message{channel: "chan1", topic: "topic", event: "join", message: %{}}

    PubSub.create("chan1:topic")
    {:ok, sockets} = Transport.dispatch(message, sockets, self, Router)
    Transport.dispatch_leave(sockets, :reason)
    assert_received :left
    assert PubSub.subscribers("chan1:topic") == []
  end

  test "#info is called when receiving regular process messages" do
    socket = %Socket{pid: self, router: Router, channel: "chan1"}
    sockets = HashDict.put(HashDict.new, {"chan1", "topic"}, socket)
    message = %Message{channel: "chan1", topic: "topic", event: "join", message: %{}}

    PubSub.create("chan1:topic")
    {:ok, sockets} = Transport.dispatch(message, sockets, self, Router)
    Transport.dispatch_info(sockets, :stuff)
    assert_received :info
  end

  test "#join raise InvalidReturn exception when return type invalid" do
    socket = %Socket{pid: self, router: Router, channel: "chan3"}
    sockets = HashDict.put(HashDict.new, {"chan3", "topic"}, socket)
    message = %Message{channel: "chan3", topic: "topic", event: "join", message: %{}}

    assert_raise InvalidReturn, fn ->
      {:ok, _sockets} = Transport.dispatch(message, sockets, self, Router)

  alias Phoenix.Channel.Transport
  alias Phoenix.Channel.Transport.InvalidReturn

  defmodule MyChannel do
    use Phoenix.Channel
    def join(_socket, _topic, msg), do: msg
    def leave(_socket, _msg) do
      Process.get(:leave)
    end
    def event(socket, "info", msg) do
      send socket.pid, :info
      msg
    end
    def event(_socket, "boom", msg), do: msg
    def event(socket, "put", dict) do
      Enum.reduce dict, socket, fn {k, v}, socket -> Socket.assign(socket, k, v) end
    end
    def event(socket, "get", %{"key" => key}) do
      send socket.pid, socket.assigns[key]
      socket
    end
  end

  defmodule Router do
    use Phoenix.Router
    use Phoenix.Router.Socket, mount: "/ws"

    channel "chan1", MyChannel
  end

  def new_socket do
    %Socket{pid: self,
            router: Router,
            topic: "topic",
            channel: "chan1",
            assigns: []}
  end
  def join_message(message) do
    %Message{channel: "chan1",
             topic: "topic",
             event: "join",
             message: message}
  end

  test "#subscribe/unsubscribe's socket to/from topic" do
