From 457cfe29661dc7625b25d0d52c1af85782d70d76 Mon Sep 17 00:00:00 2001
From: Denis Davidyuk <denis_davidyuk@hotmail.com>
Date: Sat, 22 Jul 2023 19:07:51 +0600
Subject: [PATCH 1/5] docs(aepp,wallet): fix disconnect in examples

---
 examples/browser/aepp/src/App.vue             |  2 -
 examples/browser/aepp/src/Connect.vue         | 75 ++++++++++------
 examples/browser/aepp/src/store.js            |  1 -
 examples/browser/aepp/src/styles.scss         |  2 +-
 examples/browser/wallet-iframe/src/App.vue    | 90 ++++++++++++++-----
 .../browser/wallet-iframe/src/styles.scss     |  2 +-
 .../wallet-web-extension/src/styles.scss      |  2 +-
 7 files changed, 119 insertions(+), 55 deletions(-)

diff --git a/examples/browser/aepp/src/App.vue b/examples/browser/aepp/src/App.vue
index 32bfef5240..9c3e88ffdc 100644
--- a/examples/browser/aepp/src/App.vue
+++ b/examples/browser/aepp/src/App.vue
@@ -41,7 +41,6 @@
 </template>
 
 <script>
-import { mapState } from 'vuex';
 import Connect from './Connect.vue';
 import Basic from './Basic.vue';
 import Contracts from './Contracts.vue';
@@ -53,7 +52,6 @@ export default {
     Connect, Basic, Contracts, PayForTx, TypedData,
   },
   data: () => ({ view: '' }),
-  computed: mapState(['aeSdk']),
 };
 </script>
 
diff --git a/examples/browser/aepp/src/Connect.vue b/examples/browser/aepp/src/Connect.vue
index 5ec9bcbf63..3d335e6292 100644
--- a/examples/browser/aepp/src/Connect.vue
+++ b/examples/browser/aepp/src/Connect.vue
@@ -13,18 +13,26 @@
       </label>
       <div><input v-model="reverseIframeWalletUrl"></div>
     </div>
+
+    <button
+      v-if="walletConnected"
+      @click="disconnect"
+    >
+      Disconnect
+    </button>
     <button
-      v-if="connectMethod && !walletConnected"
+      v-else-if="connectMethod"
       :disabled="walletConnecting"
       @click="connect"
     >
       Connect
     </button>
+
     <button
-      v-if="walletConnected"
-      @click="disconnect"
+      v-if="cancelWalletDetection"
+      @click="cancelWalletDetection"
     >
-      Disconnect
+      Cancel detection
     </button>
   </div>
 
@@ -34,6 +42,7 @@
       <div>
         {{
           (walletConnected && 'Wallet connected')
+          || (cancelWalletDetection && 'Wallet detection')
           || (walletConnecting && 'Wallet connecting')
           || 'Ready to connect to wallet'
         }}
@@ -48,7 +57,7 @@
 
 <script>
 import {
-  walletDetector, BrowserWindowMessageConnection, RpcConnectionDenyError,
+  walletDetector, BrowserWindowMessageConnection, RpcConnectionDenyError, RpcRejectedByUserError,
 } from '@aeternity/aepp-sdk';
 import { mapState } from 'vuex';
 
@@ -60,6 +69,7 @@ export default {
     reverseIframe: null,
     reverseIframeWalletUrl: process.env.VUE_APP_WALLET_URL ?? 'http://localhost:9000',
     walletInfo: null,
+    cancelWalletDetection: null,
   }),
   computed: {
     ...mapState(['aeSdk']),
@@ -69,32 +79,40 @@ export default {
     },
   },
   methods: {
-    async scanForWallets() {
-      return new Promise((resolve) => {
-        let stopScan;
-
-        const handleWallets = async ({ wallets, newWallet }) => {
-          newWallet = newWallet || Object.values(wallets)[0];
+    async detectWallets() {
+      if (this.connectMethod === 'reverse-iframe') {
+        this.reverseIframe = document.createElement('iframe');
+        this.reverseIframe.src = this.reverseIframeWalletUrl;
+        this.reverseIframe.style.display = 'none';
+        document.body.appendChild(this.reverseIframe);
+      }
+      const connection = new BrowserWindowMessageConnection();
+      return new Promise((resolve, reject) => {
+        const stopDetection = walletDetector(connection, async ({ newWallet }) => {
           if (confirm(`Do you want to connect to wallet ${newWallet.info.name} with id ${newWallet.info.id}`)) {
-            stopScan();
+            stopDetection();
             resolve(newWallet.getConnection());
+            this.cancelWalletDetection = null;
           }
+        });
+        this.cancelWalletDetection = () => {
+          reject(new Error('Wallet detection cancelled'));
+          stopDetection();
+          this.cancelWalletDetection = null;
+          if (this.reverseIframe) this.reverseIframe.remove();
         };
-
-        const scannerConnection = new BrowserWindowMessageConnection();
-        stopScan = walletDetector(scannerConnection, handleWallets);
       });
     },
     async connect() {
       this.walletConnecting = true;
+      this.aeSdk.onDisconnect = () => {
+        this.walletConnected = false;
+        this.walletInfo = null;
+        this.$store.commit('setAddress', undefined);
+        if (this.reverseIframe) this.reverseIframe.remove();
+      };
       try {
-        if (this.connectMethod === 'reverse-iframe') {
-          this.reverseIframe = document.createElement('iframe');
-          this.reverseIframe.src = this.reverseIframeWalletUrl;
-          this.reverseIframe.style.display = 'none';
-          document.body.appendChild(this.reverseIframe);
-        }
-        const connection = await this.scanForWallets();
+        const connection = await this.detectWallets();
         try {
           this.walletInfo = await this.aeSdk.connectToWallet(connection);
         } catch (error) {
@@ -104,14 +122,19 @@ export default {
         this.walletConnected = true;
         const { address: { current } } = await this.aeSdk.subscribeAddress('subscribe', 'connected');
         this.$store.commit('setAddress', Object.keys(current)[0]);
+      } catch (error) {
+        if (
+          error.message === 'Wallet detection cancelled'
+          || error instanceof RpcConnectionDenyError
+          || error instanceof RpcRejectedByUserError
+        ) return;
+        throw error;
       } finally {
         this.walletConnecting = false;
       }
     },
-    async disconnect() {
-      await this.aeSdk.disconnectWallet();
-      this.walletConnected = false;
-      if (this.reverseIframe) this.reverseIframe.remove();
+    disconnect() {
+      this.aeSdk.disconnectWallet();
     },
   },
 };
diff --git a/examples/browser/aepp/src/store.js b/examples/browser/aepp/src/store.js
index 32d20a5184..5205b8fd69 100644
--- a/examples/browser/aepp/src/store.js
+++ b/examples/browser/aepp/src/store.js
@@ -25,7 +25,6 @@ const store = createStore({
         store.commit('setNetworkId', networkId);
       },
       onAddressChange: ({ current }) => store.commit('setAddress', Object.keys(current)[0]),
-      onDisconnect: () => alert('Aepp is disconnected'),
     })),
   },
   mutations: {
diff --git a/examples/browser/aepp/src/styles.scss b/examples/browser/aepp/src/styles.scss
index 1ea40925ba..080d68ea8e 100644
--- a/examples/browser/aepp/src/styles.scss
+++ b/examples/browser/aepp/src/styles.scss
@@ -19,7 +19,7 @@ body {
 }
 
 button {
-  @extend .w-32, .p-2, .m-2, .rounded-full, .bg-purple-500, .text-white, .text-xs;
+  @extend .w-40, .p-2, .m-2, .rounded-full, .bg-purple-500, .text-white, .text-xs;
 
   &:disabled {
     @extend .bg-purple-300, .cursor-not-allowed;
diff --git a/examples/browser/wallet-iframe/src/App.vue b/examples/browser/wallet-iframe/src/App.vue
index b0aad366e0..e96f1a3b21 100644
--- a/examples/browser/wallet-iframe/src/App.vue
+++ b/examples/browser/wallet-iframe/src/App.vue
@@ -14,10 +14,26 @@
       <div>Balance</div>
       <Value :value="balancePromise" />
     </div>
+    <div>
+      <div>RPC client</div>
+      <div>status: {{ clientStatus ?? 'no client' }}, id: {{ clientId ?? 'not defined' }}</div>
+    </div>
 
     <button @click="switchAccount">Switch Account</button>
     <button @click="switchNode">Switch Node</button>
-    <button @click="disconnect">Disconnect</button>
+
+    <button
+      v-if="clientStatus === 'CONNECTED'"
+      @click="disconnect"
+    >
+      Disconnect
+    </button>
+    <button
+      v-else
+      @click="() => (stopSharingWalletInfo ?? shareWalletInfo)()"
+    >
+      {{ stopSharingWalletInfo ? 'Stop sharing' : 'Share wallet info' }}
+    </button>
   </div>
 
   <iframe
@@ -30,7 +46,7 @@
 <script>
 import {
   MemoryAccount, generateKeyPair, AeSdkWallet, Node, CompilerHttp,
-  BrowserWindowMessageConnection, METHODS, WALLET_TYPE,
+  BrowserWindowMessageConnection, METHODS, WALLET_TYPE, RPC_STATUS,
   RpcConnectionDenyError, RpcRejectedByUserError, unpackTx, decodeFateValue,
 } from '@aeternity/aepp-sdk';
 import Value from './Value.vue';
@@ -43,24 +59,42 @@ export default {
     nodeName: '',
     address: '',
     balancePromise: null,
+    clientId: null,
+    clientStatus: null,
+    stopSharingWalletInfo: null,
   }),
   methods: {
-    async shareWalletInfo(clientId, { interval = 5000, attemps = 5 } = {}) {
-      this.aeSdk.shareWalletInfo(clientId);
-      while (attemps) {
-        await new Promise((resolve) => {
-          setTimeout(resolve, interval);
-        });
-        this.aeSdk.shareWalletInfo(clientId);
-        attemps -= 1;
+    shareWalletInfo({ interval = 5000, attempts = 5 } = {}) {
+      const target = this.runningInFrame ? window.parent : this.$refs.aepp.contentWindow;
+      const connection = new BrowserWindowMessageConnection({ target });
+      this.clientId = this.aeSdk.addRpcClient(connection);
+
+      this.aeSdk.shareWalletInfo(this.clientId);
+      const intervalId = setInterval(() => {
+        this.aeSdk.shareWalletInfo(this.clientId);
+        attempts -= 1;
+        if (!attempts) return this.stopSharingWalletInfo();
+      }, interval);
+
+      this.stopSharingWalletInfo = () => {
+        clearInterval(intervalId);
+        // TODO: replace with clientStatus
+        const client = this.aeSdk._getClient(this.clientId);
+        if (client.status === RPC_STATUS.WAITING_FOR_CONNECTION_REQUEST) {
+          this.aeSdk.removeRpcClient(this.clientId);
+        }
+        this.stopSharingWalletInfo = null;
       }
-      console.log('Finish sharing wallet info');
     },
     disconnect() {
-      Object.values(this.aeSdk.rpcClients).forEach((client) => {
-        client.notify(METHODS.closeConnection);
-        client.disconnect();
-      });
+      // TODO: move to removeRpcClient (would be a semi-breaking change)
+      const client = this.aeSdk._getClient(this.clientId);
+      if (client.status === RPC_STATUS.CONNECTED) {
+        client.rpc.notify(METHODS.closeConnection, null);
+      }
+
+      this.aeSdk.removeRpcClient(this.clientId);
+      this.clientId = null;
     },
     async switchAccount() {
       this.address = this.aeSdk.addresses().find((a) => a !== this.address);
@@ -72,6 +106,14 @@ export default {
         .find((name) => name !== this.nodeName);
       this.aeSdk.selectNode(this.nodeName);
     },
+    updateClientStatus() {
+      if (!this.clientId) {
+        this.clientStatus = null;
+        return;
+      }
+      const client = this.aeSdk._getClient(this.clientId);
+      this.clientStatus = client.status;
+    },
   },
   mounted() {
     const aeppInfo = {};
@@ -116,7 +158,6 @@ export default {
       }
     }
 
-    let clientId;
     this.aeSdk = new AeSdkWallet({
       id: window.origin,
       type: WALLET_TYPE.window,
@@ -135,22 +176,21 @@ export default {
           throw new RpcConnectionDenyError();
         }
         aeppInfo[aeppId] = params;
+        setTimeout(() => this.stopSharingWalletInfo());
       },
       onSubscription: genConfirmCallback('subscription'),
       onAskAccounts: genConfirmCallback('get accounts'),
-      onDisconnect() {
-        this.shareWalletInfo(clientId);
+      onDisconnect: (clientId) => {
+        console.log('disconnected client', clientId);
+        this.clientId = null;
       },
     });
 
+    if (this.runningInFrame) this.shareWalletInfo();
+
     this.nodeName = this.aeSdk.selectedNodeName;
     [this.address] = this.aeSdk.addresses();
 
-    const target = this.runningInFrame ? window.parent : this.$refs.aepp.contentWindow;
-    const connection = new BrowserWindowMessageConnection({ target });
-    clientId = this.aeSdk.addRpcClient(connection);
-    this.shareWalletInfo(clientId);
-
     this.$watch(
       ({ address, nodeName }) => [address, nodeName],
       ([address]) => {
@@ -158,6 +198,10 @@ export default {
       },
       { immediate: true },
     );
+
+    // TODO: replace setInterval with subscription after refactoring
+    setInterval(() => this.updateClientStatus(), 1000);
+    this.$watch(({ clientId }) => [clientId], () => this.updateClientStatus(), { immediate: true });
   },
 };
 </script>
diff --git a/examples/browser/wallet-iframe/src/styles.scss b/examples/browser/wallet-iframe/src/styles.scss
index e595669c68..1e921cc8fb 100644
--- a/examples/browser/wallet-iframe/src/styles.scss
+++ b/examples/browser/wallet-iframe/src/styles.scss
@@ -15,7 +15,7 @@ body {
 }
 
 button {
-  @extend .w-32, .p-2, .m-2, .rounded-full, .bg-purple-500, .text-white, .text-xs;
+  @extend .w-40, .p-2, .m-2, .rounded-full, .bg-purple-500, .text-white, .text-xs;
 
   &:disabled {
     @extend .bg-purple-300, .cursor-not-allowed;
diff --git a/examples/browser/wallet-web-extension/src/styles.scss b/examples/browser/wallet-web-extension/src/styles.scss
index 28df27f2a0..6130689684 100644
--- a/examples/browser/wallet-web-extension/src/styles.scss
+++ b/examples/browser/wallet-web-extension/src/styles.scss
@@ -7,7 +7,7 @@ body {
 }
 
 button {
-  @extend .w-32, .p-2, .m-2, .rounded-full, .bg-purple-500, .text-white, .text-xs;
+  @extend .w-40, .p-2, .m-2, .rounded-full, .bg-purple-500, .text-white, .text-xs;
 
   &:disabled {
     @extend .bg-purple-300, .cursor-not-allowed;

From e2ffc25021d16b197bba41c88520e0a842b88ba9 Mon Sep 17 00:00:00 2001
From: Denis Davidyuk <denis_davidyuk@hotmail.com>
Date: Sat, 22 Jul 2023 21:53:58 +0600
Subject: [PATCH 2/5] fix(wallet): don't ask to confirm unsubscription

---
 src/AeSdkWallet.ts | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/AeSdkWallet.ts b/src/AeSdkWallet.ts
index 8a46aa61fc..879a2fd916 100644
--- a/src/AeSdkWallet.ts
+++ b/src/AeSdkWallet.ts
@@ -219,10 +219,10 @@ export default class AeSdkWallet extends AeSdk {
           [METHODS.subscribeAddress]: async ({ type, value }, origin) => {
             if (!this._isRpcClientConnected(id)) throw new RpcNotAuthorizeError();
 
-            await this.onSubscription(id, { type, value }, origin);
-
             switch (type) {
               case SUBSCRIPTION_TYPES.subscribe:
+                // TODO: remove `type` as it always subscribe
+                await this.onSubscription(id, { type, value }, origin);
                 client.addressSubscription.add(value);
                 break;
               case SUBSCRIPTION_TYPES.unsubscribe:

From 36920b4234a00be4e3cc2728372e875b361f9409 Mon Sep 17 00:00:00 2001
From: Denis Davidyuk <denis_davidyuk@hotmail.com>
Date: Sat, 22 Jul 2023 21:58:03 +0600
Subject: [PATCH 3/5] fix(wallet): don't require to be subscribed to request
 addresses

---
 src/AeSdkWallet.ts | 14 +++++---------
 1 file changed, 5 insertions(+), 9 deletions(-)

diff --git a/src/AeSdkWallet.ts b/src/AeSdkWallet.ts
index 879a2fd916..a4ac0ec0c6 100644
--- a/src/AeSdkWallet.ts
+++ b/src/AeSdkWallet.ts
@@ -109,9 +109,10 @@ export default class AeSdkWallet extends AeSdk {
   _pushAccountsToApps(): void {
     if (this._clients == null) return;
     Array.from(this._clients.keys())
-      .filter((clientId) => this._isRpcClientSubscribed(clientId))
-      .map((clientId) => this._getClient(clientId).rpc)
-      .forEach((client) => client.notify(METHODS.updateAddress, this.getAccounts()));
+      .filter((clientId) => this._isRpcClientConnected(clientId))
+      .map((clientId) => this._getClient(clientId))
+      .filter((client) => client.addressSubscription.size !== 0)
+      .forEach((client) => client.rpc.notify(METHODS.updateAddress, this.getAccounts()));
   }
 
   override selectAccount(address: Encoded.AccountAddress): void {
@@ -149,11 +150,6 @@ export default class AeSdkWallet extends AeSdk {
     return client;
   }
 
-  _isRpcClientSubscribed(clientId: string): boolean {
-    return this._isRpcClientConnected(clientId)
-      && this._getClient(clientId).addressSubscription.size !== 0;
-  }
-
   _isRpcClientConnected(clientId: string): boolean {
     return RPC_STATUS.CONNECTED === this._getClient(clientId).status
       && this._getClient(clientId).rpc.connection.isConnected();
@@ -238,7 +234,7 @@ export default class AeSdkWallet extends AeSdk {
             };
           },
           [METHODS.address]: async (params, origin) => {
-            if (!this._isRpcClientSubscribed(id)) throw new RpcNotAuthorizeError();
+            if (!this._isRpcClientConnected(id)) throw new RpcNotAuthorizeError();
             await this.onAskAccounts(id, params, origin);
             return this.addresses();
           },

From fa900c0481cc898211a3d99d22d57d34cac02177 Mon Sep 17 00:00:00 2001
From: Denis Davidyuk <denis_davidyuk@hotmail.com>
Date: Sat, 22 Jul 2023 22:03:28 +0600
Subject: [PATCH 4/5] fix(wallet): return accounts according to subscription

---
 src/AeSdkWallet.ts | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/src/AeSdkWallet.ts b/src/AeSdkWallet.ts
index a4ac0ec0c6..a0217d54ac 100644
--- a/src/AeSdkWallet.ts
+++ b/src/AeSdkWallet.ts
@@ -106,13 +106,23 @@ export default class AeSdkWallet extends AeSdk {
     this._type = type;
   }
 
+  _getAccountsForClient({ addressSubscription }: RpcClientsInfo): Accounts {
+    const { current, connected } = this.getAccounts();
+    return {
+      current: addressSubscription.has('current') || addressSubscription.has('connected')
+        ? current : {},
+      connected: addressSubscription.has('connected') ? connected : {},
+    };
+  }
+
   _pushAccountsToApps(): void {
     if (this._clients == null) return;
     Array.from(this._clients.keys())
       .filter((clientId) => this._isRpcClientConnected(clientId))
       .map((clientId) => this._getClient(clientId))
       .filter((client) => client.addressSubscription.size !== 0)
-      .forEach((client) => client.rpc.notify(METHODS.updateAddress, this.getAccounts()));
+      .forEach((client) => client.rpc
+        .notify(METHODS.updateAddress, this._getAccountsForClient(client)));
   }
 
   override selectAccount(address: Encoded.AccountAddress): void {
@@ -230,7 +240,7 @@ export default class AeSdkWallet extends AeSdk {
 
             return {
               subscription: Array.from(client.addressSubscription),
-              address: this.getAccounts(),
+              address: this._getAccountsForClient(client),
             };
           },
           [METHODS.address]: async (params, origin) => {

From 1d848437498a60d36532abd1a8ab51613e35b99a Mon Sep 17 00:00:00 2001
From: Denis Davidyuk <denis_davidyuk@hotmail.com>
Date: Sun, 23 Jul 2023 16:08:10 +0600
Subject: [PATCH 5/5] test(aepp,wallet): use `assertNotNull` instead throwing
 `UnexpectedTsError`

---
 test/integration/rpc.ts | 20 ++++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/test/integration/rpc.ts b/test/integration/rpc.ts
index 65f631bb39..4659dbbf6c 100644
--- a/test/integration/rpc.ts
+++ b/test/integration/rpc.ts
@@ -24,7 +24,6 @@ import {
   verify,
   NoWalletConnectedError,
   UnAuthorizedAccountError,
-  UnexpectedTsError,
   UnknownRpcClientError,
   UnsubscribedAccountError,
   AccountBase,
@@ -38,6 +37,7 @@ import {
   getSdk, ignoreVersion, networkId, url, compilerUrl,
 } from '.';
 import { Accounts, Network } from '../../src/aepp-wallet-communication/rpc/types';
+import { assertNotNull } from '../utils';
 
 const WindowPostMessageFake = (
   name: string,
@@ -149,7 +149,6 @@ describe('Aepp<->Wallet', function aeppWallet() {
 
     it('Should receive `announcePresence` message from wallet', async () => {
       const isReceived = new Promise<boolean>((resolve) => {
-        if (connections.aeppWindow.addEventListener == null) throw new UnexpectedTsError();
         connections.aeppWindow.addEventListener('message', (msg) => {
           resolve(msg.data.method === 'connection.announcePresence');
         });
@@ -321,7 +320,7 @@ describe('Aepp<->Wallet', function aeppWallet() {
 
     it('Try to sign using unpermited account', async () => {
       const { publicKey: pub } = generateKeyPair();
-      if (aepp.rpcClient == null) throw new UnexpectedTsError();
+      assertNotNull(aepp.rpcClient);
       await expect(aepp.rpcClient.request(METHODS.sign, {
         tx: 'tx_+NkLAfhCuECIIeWttRUiZ32uriBdmM1t+dCg90KuG2ABxOiuXqzpAul6uTWvsyfx3EFJDah6trudrityh+6XSX3mkPEimhgGuJH4jzIBoQELtO15J/l7UeG8teE0DRIzWyorEsi8UiHWPEvLOdQeYYgbwW1nTsgAAKEB6bv2BOYRtUYKOzmZ6Xcbb2BBfXPOfFUZ4S9+EnoSJcqIG8FtZ07IAACIAWNFeF2KAAAKAIYSMJzlQADAoDBrIcoop8JfZ4HOD9p3nDTiNthj7jjl+ArdHwEMUrvQgitwOr/v3Q==',
         onAccount: pub,
@@ -351,8 +350,9 @@ describe('Aepp<->Wallet', function aeppWallet() {
         payload: encode(Buffer.from('zerospend'), Encoding.Bytearray),
       });
       const res = await aepp.sendTransaction(tx);
-      if (res.tx?.payload == null || res.blockHeight == null) throw new UnexpectedTsError();
+      assertNotNull(res.tx?.payload);
       decode(res.tx.payload as Encoded.Any).toString().should.be.equal('zerospend2');
+      assertNotNull(res.blockHeight);
       res.blockHeight.should.be.a('number');
     });
 
@@ -454,7 +454,7 @@ describe('Aepp<->Wallet', function aeppWallet() {
 
     describe('Subscriptions', () => {
       it('Add new account to wallet: receive notification for update accounts', async () => {
-        if (aepp._accounts == null) throw new UnexpectedTsError();
+        assertNotNull(aepp._accounts);
         const connectedLength = Object.keys(aepp._accounts.connected).length;
         const accountsPromise = new Promise<Accounts>((resolve) => {
           aepp.onAddressChange = resolve;
@@ -464,7 +464,7 @@ describe('Aepp<->Wallet', function aeppWallet() {
       });
 
       it('Receive update for wallet select account', async () => {
-        if (aepp._accounts == null) throw new UnexpectedTsError();
+        assertNotNull(aepp._accounts);
         const connectedAccount = Object
           .keys(aepp._accounts.connected)[0] as Encoded.AccountAddress;
         const accountsPromise = new Promise<Accounts>((resolve) => {
@@ -472,7 +472,6 @@ describe('Aepp<->Wallet', function aeppWallet() {
         });
         wallet.selectAccount(connectedAccount);
         const { connected, current } = await accountsPromise;
-        if (current == null || connected == null) throw new UnexpectedTsError();
         expect(current[connectedAccount]).to.be.eql({});
         expect(Object.keys(connected).includes(connectedAccount)).to.be.equal(false);
       });
@@ -488,7 +487,7 @@ describe('Aepp<->Wallet', function aeppWallet() {
     });
 
     it('Try to connect unsupported protocol', async () => {
-      if (aepp.rpcClient == null) throw new UnexpectedTsError();
+      assertNotNull(aepp.rpcClient);
       await expect(aepp.rpcClient.request(METHODS.connect, { name: 'test-aepp', version: 2 as 1, connectNode: false })).to.be.eventually.rejectedWith('Unsupported Protocol Version').with.property('code', 5);
     });
 
@@ -610,8 +609,9 @@ describe('Aepp<->Wallet', function aeppWallet() {
         payload: encode(Buffer.from('zerospend'), Encoding.Bytearray),
       });
       const res = await aepp.sendTransaction(tx);
-      if (res.tx?.payload == null || res.blockHeight == null) throw new UnexpectedTsError();
+      assertNotNull(res.tx?.payload);
       decode(res.tx.payload as Encoded.Any).toString().should.be.equal('zerospend2');
+      assertNotNull(res.blockHeight);
       res.blockHeight.should.be.a('number');
     });
 
@@ -621,7 +621,7 @@ describe('Aepp<->Wallet', function aeppWallet() {
         wallet.addNode('second_node', node, true);
       });
       message.networkId.should.be.equal(networkId);
-      if (message.node == null) throw new UnexpectedTsError();
+      assertNotNull(message.node);
       message.node.should.be.an('object');
       expect(wallet.selectedNodeName).to.be.equal('second_node');
     });