1 files changed, 12 insertions, 8 deletions
diff --git a/actioncable/test/channel/stream_test.rb b/actioncable/test/channel/stream_test.rb
index df9d44d8dd..53dd7e5b77 100644
--- a/actioncable/test/channel/stream_test.rb
+++ b/actioncable/test/channel/stream_test.rb
@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require "test_helper"
+require "active_support/testing/method_call_assertions"
 require "stubs/test_connection"
 require "stubs/room"
 
@@ -143,6 +144,8 @@ module ActionCable::StreamTests
   end
 
   class StreamFromTest < ActionCable::TestCase
+    include ActiveSupport::Testing::MethodCallAssertions
+
     setup do
       @server = TestServer.new(subscription_adapter: ActionCable::SubscriptionAdapter::Async)
       @server.config.allowed_request_origins = %w( http://rubyonrails.com )
@@ -153,10 +156,11 @@ module ActionCable::StreamTests
         connection = open_connection
         subscribe_to connection, identifiers: { id: 1 }
 
-        connection.websocket.expects(:transmit)
-        @server.broadcast "test_room_1", { foo: "bar" }, { coder: DummyEncoder }
-        wait_for_async
-        wait_for_executor connection.server.worker_pool.executor
+        assert_called(connection.websocket, :transmit) do
+          @server.broadcast "test_room_1", { foo: "bar" }, { coder: DummyEncoder }
+          wait_for_async
+          wait_for_executor connection.server.worker_pool.executor
+        end
       end
     end
 
@@ -175,10 +179,10 @@ module ActionCable::StreamTests
       run_in_eventmachine do
         connection = open_connection
         expected = { "identifier" => { "channel" => MultiChatChannel.name }.to_json, "type" => "confirm_subscription" }
-        connection.websocket.expects(:transmit).with(expected.to_json)
-        receive(connection, command: "subscribe", channel: MultiChatChannel.name, identifiers: {})
-
-        wait_for_async
+        assert_called(connection.websocket, :transmit, [expected.to_json]) do
+          receive(connection, command: "subscribe", channel: MultiChatChannel.name, identifiers: {})
+          wait_for_async
+        end
       end
     end