From 469e9d6bf3ea93b859c172cf9883ad2502d4990c Mon Sep 17 00:00:00 2001
From: Szymon Lesisz <sz.lesisz@gmail.com>
Date: Thu, 4 Oct 2018 10:48:23 +0200
Subject: [PATCH] move all DiscoveryActions from TrezorConnectService to
 WalletService

---
 src/services/TrezorConnectService.js |  6 ------
 src/services/WalletService.js        | 19 +++++++++++++++++++
 2 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/src/services/TrezorConnectService.js b/src/services/TrezorConnectService.js
index c14f0a24..397f9e0d 100644
--- a/src/services/TrezorConnectService.js
+++ b/src/services/TrezorConnectService.js
@@ -4,7 +4,6 @@ import {
     TRANSPORT, DEVICE, BLOCKCHAIN,
 } from 'trezor-connect';
 import * as TrezorConnectActions from 'actions/TrezorConnectActions';
-import * as DiscoveryActions from 'actions/DiscoveryActions';
 import * as BlockchainActions from 'actions/BlockchainActions';
 import * as RouterActions from 'actions/RouterActions';
 import * as ModalActions from 'actions/ModalActions';
@@ -55,14 +54,9 @@ const TrezorConnectService: Middleware = (api: MiddlewareAPI) => (next: Middlewa
             api.dispatch(RouterActions.selectFirstAvailableDevice());
         }
     } else if (action.type === DEVICE.CONNECT || action.type === DEVICE.CONNECT_UNACQUIRED) {
-        api.dispatch(DiscoveryActions.restore());
         api.dispatch(ModalActions.onDeviceConnect(action.device));
-    } else if (action.type === CONNECT.AUTH_DEVICE) {
-        api.dispatch(DiscoveryActions.check());
     } else if (action.type === CONNECT.DUPLICATE) {
         api.dispatch(RouterActions.selectDevice(action.device));
-    } else if (action.type === CONNECT.COIN_CHANGED) {
-        api.dispatch(TrezorConnectActions.coinChanged(action.payload.network));
     } else if (action.type === BLOCKCHAIN.BLOCK) {
         api.dispatch(BlockchainActions.onBlockMined(action.payload.coin));
     } else if (action.type === BLOCKCHAIN.NOTIFICATION) {
diff --git a/src/services/WalletService.js b/src/services/WalletService.js
index 76db7638..3cfe4ddb 100644
--- a/src/services/WalletService.js
+++ b/src/services/WalletService.js
@@ -10,6 +10,7 @@ import * as LocalStorageActions from 'actions/LocalStorageActions';
 import * as TrezorConnectActions from 'actions/TrezorConnectActions';
 import * as SelectedAccountActions from 'actions/SelectedAccountActions';
 import * as SendFormActionActions from 'actions/SendFormActions';
+import * as DiscoveryActions from 'actions/DiscoveryActions';
 
 import type {
     Middleware,
@@ -77,6 +78,12 @@ const WalletService: Middleware = (api: MiddlewareAPI) => (next: MiddlewareDispa
                     network: currentLocation.state.network,
                 },
             });
+
+            // try to stop currently running discovery on previous network
+            api.dispatch(DiscoveryActions.stop());
+
+            // try to start new discovery on new network
+            api.dispatch(DiscoveryActions.restore());
         }
 
         // watch for account change
@@ -95,6 +102,18 @@ const WalletService: Middleware = (api: MiddlewareAPI) => (next: MiddlewareDispa
             // if "selectedAccount" didn't change observe send form props changes
             api.dispatch(SendFormActionActions.observe(prevState, action));
         }
+    } else if (action.type === CONNECT.AUTH_DEVICE) {
+        // selected device did changed
+        // try to restore discovery after device authentication
+        api.dispatch(DiscoveryActions.restore());
+    }
+
+    // even if "selectedDevice" didn't change because it was updated on DEVICE.CHANGED before DEVICE.CONNECT action
+    // try to restore discovery
+    if (action.type === DEVICE.CONNECT) {
+        api.dispatch(DiscoveryActions.restore());
+    } else if (action.type === DEVICE.DISCONNECT) {
+        api.dispatch(DiscoveryActions.stop());
     }
 
     return action;