2 files changed, 18 insertions, 7 deletions
diff --git a/test/components/Connect.spec.js b/test/components/Connect.spec.js
index 2a385fed5f..30228fc4c6 100644
--- a/test/components/Connect.spec.js
+++ b/test/components/Connect.spec.js
@@ -94,7 +94,7 @@ describe('components/Connect', () => {
 
   it('invokes the onConnect prop', (done) => {
     const component = renderNotConnected({
-      onConnect: done,
+      onConnect: () => done(),
     });
     const connectButton = component.find('.button .button--positive');
 
diff --git a/test/connect.spec.js b/test/connect.spec.js
index 98aba67356..bb6ac41526 100644
--- a/test/connect.spec.js
+++ b/test/connect.spec.js
@@ -13,7 +13,9 @@ describe('connect', () => {
     const chain = new IpcChain(mockIpc);
     chain.require('setCustomRelay')
       .withInputValidation(
-        (relayHostName) => expect(relayHostName).to.equal('example.com')
+        (relayEndpoint) => {
+          expect(relayEndpoint).to.equal(arbitraryRelay);
+        },
       )
       .done();
 
@@ -22,7 +24,7 @@ describe('connect', () => {
 
     chain.onSuccessOrFailure(done);
 
-    store.dispatch(connectionActions.connect(backend, 'example.com'));
+    store.dispatch(connectionActions.connect(backend, arbitraryRelay));
   });
 
   it('should set the connection state to \'disconnected\' on failed attempts', (done) => {
@@ -35,7 +37,7 @@ describe('connect', () => {
 
     expect(store.getState().connection.status).not.to.equal('disconnected');
 
-    store.dispatch(connectionActions.connect(backend, 'example.com'));
+    store.dispatch(connectionActions.connect(backend, arbitraryRelay));
 
 
     checkNextTick(() => {
@@ -45,13 +47,17 @@ describe('connect', () => {
 
   it('should update the state with the server address', () => {
     const { store, backend } = setupBackendAndStore();
-    const arbitraryString = 'www.example.com';
+    const relay = {
+      host: 'www.example.com',
+      port: 1,
+      protocol: 'udp',
+    };
 
-    return backend.connect(arbitraryString)
+    return backend.connect(relay)
       .then( () => {
         const state = store.getState().connection;
         expect(state.status).to.equal('connecting');
-        expect(state.serverAddress).to.equal(arbitraryString);
+        expect(state.serverAddress).to.equal('www.example.com');
       });
   });
 
@@ -81,3 +87,8 @@ describe('connect', () => {
   });
 });
 
+const arbitraryRelay = {
+  host: 'www.example.com',
+  port: 1,
+  protocol: 'udp',
+};