diff --git a/services/real-time/test/unit/coffee/DocumentUpdaterControllerTests.coffee b/services/real-time/test/unit/coffee/DocumentUpdaterControllerTests.coffee
index 24551396d9..b5574c8d16 100644
--- a/services/real-time/test/unit/coffee/DocumentUpdaterControllerTests.coffee
+++ b/services/real-time/test/unit/coffee/DocumentUpdaterControllerTests.coffee
@@ -11,6 +11,7 @@ describe "DocumentUpdaterController", ->
 		@callback = sinon.stub()
 		@io = { "mock": "socket.io" }
 		@rclient = []
+		@RoomEvents = { on: sinon.stub() }
 		@EditorUpdatesController = SandboxedModule.require modulePath, requires:
 			"logger-sharelatex": @logger = { error: sinon.stub(), log: sinon.stub(), warn: sinon.stub() }
 			"settings-sharelatex": @settings =
@@ -28,6 +29,8 @@ describe "DocumentUpdaterController", ->
 			"./EventLogger": @EventLogger = {checkEventOrder: sinon.stub()}
 			"./HealthCheckManager": {check: sinon.stub()}
 			"metrics-sharelatex": @metrics = {inc: sinon.stub()}
+			"./RoomManager" : @RoomManager = { eventSource: sinon.stub().returns @RoomEvents}
+			"./ChannelManager": @ChannelManager = {}
 
 	describe "listenForUpdatesFromDocumentUpdater", ->
 		beforeEach ->
diff --git a/services/real-time/test/unit/coffee/WebsocketControllerTests.coffee b/services/real-time/test/unit/coffee/WebsocketControllerTests.coffee
index 1ddf4e6359..ab442006c2 100644
--- a/services/real-time/test/unit/coffee/WebsocketControllerTests.coffee
+++ b/services/real-time/test/unit/coffee/WebsocketControllerTests.coffee
@@ -36,7 +36,7 @@ describe 'WebsocketController', ->
 			"metrics-sharelatex": @metrics =
 				inc: sinon.stub()
 				set: sinon.stub()
-				
+			"./RoomManager": @RoomManager = {}
 	
 	afterEach ->
 		tk.reset()
@@ -54,6 +54,7 @@ describe 'WebsocketController', ->
 				@privilegeLevel = "owner"
 				@ConnectedUsersManager.updateUserPosition = sinon.stub().callsArg(4)
 				@WebApiManager.joinProject = sinon.stub().callsArgWith(2, null, @project, @privilegeLevel)
+				@RoomManager.joinProject = sinon.stub()
 				@WebsocketController.joinProject @client, @user, @project_id, @callback
 				
 			it "should load the project from web", ->
@@ -62,7 +63,7 @@ describe 'WebsocketController', ->
 					.should.equal true
 					
 			it "should join the project room", ->
-				@client.join.calledWith(@project_id).should.equal true
+				@RoomManager.joinProject.calledWith(@client, @project_id).should.equal true
 					
 			it "should set the privilege level on the client", ->
 				@client.set.calledWith("privilege_level", @privilegeLevel).should.equal true
@@ -125,6 +126,7 @@ describe 'WebsocketController', ->
 			@DocumentUpdaterManager.flushProjectToMongoAndDelete = sinon.stub().callsArg(1)
 			@ConnectedUsersManager.markUserAsDisconnected = sinon.stub().callsArg(2)
 			@WebsocketLoadBalancer.emitToRoom = sinon.stub()
+			@RoomManager.leaveProjectAndDocs = sinon.stub()
 			@clientsInRoom = []
 			@io =
 				sockets:
@@ -160,6 +162,11 @@ describe 'WebsocketController', ->
 			it "should increment the leave-project metric", ->
 				@metrics.inc.calledWith("editor.leave-project").should.equal true
 			
+			it "should track the disconnection in RoomManager", ->
+				@RoomManager.leaveProjectAndDocs
+					.calledWith(@client)
+					.should.equal true
+
 		describe "when the project is not empty", ->
 			beforeEach ->
 				@clientsInRoom = ["mock-remaining-client"]
@@ -230,6 +237,7 @@ describe 'WebsocketController', ->
 			@AuthorizationManager.addAccessToDoc = sinon.stub()
 			@AuthorizationManager.assertClientCanViewProject = sinon.stub().callsArgWith(1, null)
 			@DocumentUpdaterManager.getDocument = sinon.stub().callsArgWith(3, null, @doc_lines, @version, @ranges, @ops)
+			@RoomManager.joinDoc = sinon.stub()
 
 		describe "works", ->
 			beforeEach ->
@@ -251,8 +259,8 @@ describe 'WebsocketController', ->
 					.should.equal true
 
 			it "should join the client to room for the doc_id", ->
-				@client.join
-					.calledWith(@doc_id)
+				@RoomManager.joinDoc
+					.calledWith(@client, @doc_id)
 					.should.equal true
 
 			it "should call the callback with the lines, version, ranges and ops", ->
@@ -330,11 +338,12 @@ describe 'WebsocketController', ->
 		beforeEach ->
 			@doc_id = "doc-id-123"			
 			@client.params.project_id = @project_id
+			@RoomManager.leaveDoc = sinon.stub()
 			@WebsocketController.leaveDoc @client, @doc_id, @callback
 			
 		it "should remove the client from the doc_id room", ->
-			@client.leave
-				.calledWith(@doc_id).should.equal true
+			@RoomManager.leaveDoc
+				.calledWith(@client, @doc_id).should.equal true
 				
 		it "should call the callback", ->
 			@callback.called.should.equal true
diff --git a/services/real-time/test/unit/coffee/WebsocketLoadBalancerTests.coffee b/services/real-time/test/unit/coffee/WebsocketLoadBalancerTests.coffee
index 14df2df851..c4f4519790 100644
--- a/services/real-time/test/unit/coffee/WebsocketLoadBalancerTests.coffee
+++ b/services/real-time/test/unit/coffee/WebsocketLoadBalancerTests.coffee
@@ -6,6 +6,7 @@ modulePath = require('path').join __dirname, '../../../app/js/WebsocketLoadBalan
 describe "WebsocketLoadBalancer", ->
 	beforeEach ->
 		@rclient = {}
+		@RoomEvents = {on: sinon.stub()}
 		@WebsocketLoadBalancer = SandboxedModule.require modulePath, requires:
 			"./RedisClientManager": 
 				createClientList: () => []
@@ -14,6 +15,8 @@ describe "WebsocketLoadBalancer", ->
 				parse: (data, cb) => cb null, JSON.parse(data)
 			"./EventLogger": {checkEventOrder: sinon.stub()}
 			"./HealthCheckManager": {check: sinon.stub()}
+			"./RoomManager" : @RoomManager = {eventSource: sinon.stub().returns @RoomEvents}
+			"./ChannelManager": @ChannelManager = {publish: sinon.stub()}
 		@io = {}
 		@WebsocketLoadBalancer.rclientPubList = [{publish: sinon.stub()}]
 		@WebsocketLoadBalancer.rclientSubList = [{
@@ -30,8 +33,8 @@ describe "WebsocketLoadBalancer", ->
 			@WebsocketLoadBalancer.emitToRoom(@room_id, @message, @payload...)
 
 		it "should publish the message to redis", ->
-			@WebsocketLoadBalancer.rclientPubList[0].publish
-				.calledWith("editor-events", JSON.stringify(
+			@ChannelManager.publish
+				.calledWith(@WebsocketLoadBalancer.rclientPubList[0], "editor-events", @room_id, JSON.stringify(
 					room_id: @room_id,
 					message: @message
 					payload: @payload