diff --git a/websocket-gateway/spec/jam_websockets/router_spec.rb b/websocket-gateway/spec/jam_websockets/router_spec.rb
index 01dc92549..75403152b 100644
--- a/websocket-gateway/spec/jam_websockets/router_spec.rb
+++ b/websocket-gateway/spec/jam_websockets/router_spec.rb
@@ -51,7 +51,7 @@ def login(router, user, password, client_id, token, client_type)
     args[0].should == client
     args[1].is_a?(Jampb::ClientMessage).should be_true
   end
-  router.should_receive(:extract_ip).at_least(:once).with(client).and_return("127.0.0.1")
+  router.should_receive(:extract_ip).at_least(:once).with(client, nil).and_return("127.0.0.1")
   client.should_receive(:onclose)
   client.should_receive(:onerror)
 
@@ -61,6 +61,7 @@ def login(router, user, password, client_id, token, client_type)
   @router.new_client(client, false)
   handshake = double("handshake")
   handshake.should_receive(:query).exactly(3).times.and_return({ "pb" => "true", "channel_id" => SecureRandom.uuid, "client_id" => client_id, "token" => token, "client_type" => client_type })
+  handshake.should_receive(:headers).at_least(:once).and_return({})
   client.onopenblock.call handshake
 
   client
@@ -84,7 +85,7 @@ def login_latency_tester(router, latency_tester, client_id)
     args[0].should == client
     args[1].is_a?(Jampb::ClientMessage).should be_true
   end
-  router.should_receive(:extract_ip).at_least(:once).with(client).and_return("127.0.0.1")
+  router.should_receive(:extract_ip).at_least(:once).with(client, nil).and_return("127.0.0.1")
   client.should_receive(:onclose)
   client.should_receive(:onerror)
 
@@ -93,6 +94,7 @@ def login_latency_tester(router, latency_tester, client_id)
   @router.new_client(client, true)
   handshake = double("handshake")
   handshake.should_receive(:query).exactly(3).times.and_return({ "pb" => "true", "channel_id" => SecureRandom.uuid, "client_type" => "latency_tester", "client_id" => client_id })
+  handshake.should_receive(:headers).at_least(:once).and_return({})
   client.onopenblock.call handshake
 
   client