Reorganize `MockWebSocket`

2 files changed, 32 insertions, 22 deletions

diff --git a/actioncable/test/javascript/src/test_helpers/mock_websocket.coffee b/actioncable/test/javascript/src/test_helpers/mock_websocket.coffee
index 6508984c37..147a336efb 100644
--- a/actioncable/test/javascript/src/test_helpers/mock_websocket.coffee
+++ b/actioncable/test/javascript/src/test_helpers/mock_websocket.coffee
@@ -1,10 +1,11 @@
 #= require mock-socket
 
+ActionCable.TestHelpers.MockWebSocket = MockWebSocket
+
 server = null
 consumer = null
 
 ActionCable.TestHelpers.createConsumer = (url, callback) ->
-  ActionCable.WebSocket = MockWebSocket
   server = new MockServer url
   consumer = ActionCable.createConsumer(url)
   callback(consumer, server)
diff --git a/actioncable/test/javascript/src/unit/consumer_test.coffee b/actioncable/test/javascript/src/unit/consumer_test.coffee
index d8b1450ad8..fd76b086ed 100644
--- a/actioncable/test/javascript/src/unit/consumer_test.coffee
+++ b/actioncable/test/javascript/src/unit/consumer_test.coffee
@@ -1,31 +1,40 @@
 {module, test} = QUnit
-{testURL, createConsumer} = ActionCable.TestHelpers
+{testURL, MockWebSocket, createConsumer} = ActionCable.TestHelpers
 
 module "ActionCable.Consumer", ->
-  test "#connect", (assert) ->
-    done = assert.async()
+  module 'with MockWebSocket', (hooks) ->
+    originalWebSocket = ActionCable.WebSocket
 
-    createConsumer testURL, (consumer, server) ->
-      server.on "connection", ->
-        clients = server.clients()
-        assert.equal clients.length, 1
-        assert.equal clients[0].readyState, WebSocket.OPEN
-        done()
+    hooks.beforeEach ->
+      ActionCable.WebSocket = MockWebSocket
 
-      consumer.connect()
+    hooks.afterEach ->
+      ActionCable.WebSocket = originalWebSocket
 
-  test "#disconnect", (assert) ->
-    done = assert.async()
+    test "#connect", (assert) ->
+      done = assert.async()
 
-    createConsumer testURL, (consumer, server) ->
-      server.on "connection", ->
-        clients = server.clients()
-        assert.equal clients.length, 1
-
-        clients[0].addEventListener "close", (event) ->
-          assert.equal event.type, "close"
+      createConsumer testURL, (consumer, server) ->
+        server.on "connection", ->
+          clients = server.clients()
+          assert.equal clients.length, 1
+          assert.equal clients[0].readyState, WebSocket.OPEN
           done()
 
-        consumer.disconnect()
+        consumer.connect()
+
+    test "#disconnect", (assert) ->
+      done = assert.async()
+
+      createConsumer testURL, (consumer, server) ->
+        server.on "connection", ->
+          clients = server.clients()
+          assert.equal clients.length, 1
+
+          clients[0].addEventListener "close", (event) ->
+            assert.equal event.type, "close"
+            done()
+
+          consumer.disconnect()
 
-      consumer.connect()
+        consumer.connect()