4 files changed, 36 insertions, 22 deletions
diff --git a/lib/action_cable/connection.rb b/lib/action_cable/connection.rb
index 09ef1699a6..1b4a6ecc23 100644
--- a/lib/action_cable/connection.rb
+++ b/lib/action_cable/connection.rb
@@ -5,6 +5,7 @@ module ActionCable
     autoload :Identification, 'action_cable/connection/identification'
     autoload :InternalChannel, 'action_cable/connection/internal_channel'
     autoload :MessageBuffer, 'action_cable/connection/message_buffer'
+    autoload :WebSocket, 'action_cable/connection/web_socket'
     autoload :Subscriptions, 'action_cable/connection/subscriptions'
     autoload :TaggedLoggerProxy, 'action_cable/connection/tagged_logger_proxy'
   end
diff --git a/lib/action_cable/connection/base.rb b/lib/action_cable/connection/base.rb
index ae9dd58ab4..efabe40b73 100644
--- a/lib/action_cable/connection/base.rb
+++ b/lib/action_cable/connection/base.rb
@@ -14,6 +14,7 @@ module ActionCable
 
         @logger = initialize_tagged_logger
 
+        @websocket      = ActionCable::Connection::WebSocket.new(env)
         @heartbeat      = ActionCable::Connection::Heartbeat.new(self)
         @subscriptions  = ActionCable::Connection::Subscriptions.new(self)
         @message_buffer = ActionCable::Connection::MessageBuffer.new(self)
@@ -21,12 +22,10 @@ module ActionCable
         @started_at = Time.now
       end
 
-      def process
+      def response
         logger.info started_request_message
 
-        if websocket_request?
-          websocket_initialization
-
+        if websocket.possible?
           websocket.on(:open)    { |event| send_async :on_open   }
           websocket.on(:message) { |event| on_message event.data }
           websocket.on(:close)   { |event| send_async :on_close  }
@@ -38,7 +37,7 @@ module ActionCable
       end
 
       def receive(data_in_json)
-        if websocket_alive?
+        if websocket.alive?
           subscriptions.execute_command ActiveSupport::JSON.decode(data_in_json)
         else
           logger.error "Received data without a live websocket (#{data.inspect})"
@@ -46,7 +45,7 @@ module ActionCable
       end
 
       def transmit(data)
-        websocket.send data
+        websocket.transmit data
       end
 
       def close
@@ -78,19 +77,6 @@ module ActionCable
         attr_reader :websocket
         attr_reader :heartbeat, :subscriptions, :message_buffer
 
-        def websocket_initialization
-          @websocket = Faye::WebSocket.new(@env)
-        end
-
-        def websocket_alive?
-          websocket && websocket.ready_state == Faye::WebSocket::API::OPEN
-        end
-
-        def websocket_request?
-          @is_websocket ||= Faye::WebSocket.websocket?(@env)
-        end
-
-
         def on_open
           server.add_connection(self)
 
@@ -134,7 +120,7 @@ module ActionCable
           'Started %s "%s"%s for %s at %s' % [
             request.request_method,
             request.filtered_path,
-            websocket_request? ? ' [Websocket]' : '',
+            websocket.possible? ? ' [Websocket]' : '',
             request.ip,
             Time.now.to_default_s ]
         end
@@ -142,7 +128,7 @@ module ActionCable
         def finished_request_message
           'Finished "%s"%s for %s at %s' % [
             request.filtered_path,
-            websocket_request? ? ' [Websocket]' : '',
+            websocket.possible? ? ' [Websocket]' : '',
             request.ip,
             Time.now.to_default_s ]
         end
diff --git a/lib/action_cable/connection/web_socket.rb b/lib/action_cable/connection/web_socket.rb
new file mode 100644
index 0000000000..135a28cfe4
--- /dev/null
+++ b/lib/action_cable/connection/web_socket.rb
@@ -0,0 +1,27 @@
+module ActionCable
+  module Connection
+    # Decorate the Faye::WebSocket with helpers we need.
+    class WebSocket
+      delegate :rack_response, :close, :on, to: :websocket
+
+      def initialize(env)
+        @websocket = Faye::WebSocket.websocket?(env) ? Faye::WebSocket.new(env) : nil
+      end
+
+      def possible?
+        websocket
+      end
+
+      def alive?
+        websocket && websocket.ready_state == Faye::WebSocket::API::OPEN
+      end
+
+      def transmit(data)
+        websocket.send data
+      end
+
+      private
+        attr_reader :websocket
+    end
+  end
+end
diff --git a/lib/action_cable/server.rb b/lib/action_cable/server.rb
index 3a16f51757..dbfadaa34c 100644
--- a/lib/action_cable/server.rb
+++ b/lib/action_cable/server.rb
@@ -17,7 +17,7 @@ module ActionCable
     end
 
     def call(env)
-      @connection_class.new(self, env).process
+      @connection_class.new(self, env).response
     end
 
     def worker_pool