From 20c3023b9405f3358a9038bc36da7d3d92118731 Mon Sep 17 00:00:00 2001
From: Bruno Windels <brunow@matrix.org>
Date: Tue, 16 Apr 2019 11:55:33 +0200
Subject: [PATCH] use session.poll as well for polling when receiving a message

---
 src/session.js           |  9 ++++-----
 src/usecases/signup.js   |  2 +-
 src/usecases/timeline.js | 22 +++++++++-------------
 3 files changed, 14 insertions(+), 19 deletions(-)

diff --git a/src/session.js b/src/session.js
index 126112d7c7..e1d0daf7fe 100644
--- a/src/session.js
+++ b/src/session.js
@@ -202,13 +202,12 @@ module.exports = class RiotSession {
         return this.browser.close();
     }
 
-    async poll(callback, timeout) {
-        const INTERVAL = 100;
+    async poll(callback, timeout, interval = 100) {
         let waited = 0;
         while(waited < timeout) {
-            await this.delay(INTERVAL);
-            waited += INTERVAL;
-            if (callback()) {
+            await this.delay(interval);
+            waited += interval;
+            if (await callback()) {
                 return true;
             }
         }
diff --git a/src/usecases/signup.js b/src/usecases/signup.js
index 1f9605c30e..17b9669e6c 100644
--- a/src/usecases/signup.js
+++ b/src/usecases/signup.js
@@ -65,7 +65,7 @@ module.exports = async function signup(session, username, password, homeserver)
     //wait for registration to finish so the hash gets set
     //onhashchange better?
 
-    const foundHomeUrl = await session.poll(() => {
+    const foundHomeUrl = await session.poll(async () => {
         const url = session.page.url();
         return url === session.url('/#/home');
     }, 5000);
diff --git a/src/usecases/timeline.js b/src/usecases/timeline.js
index 5556660e1c..7b28328e49 100644
--- a/src/usecases/timeline.js
+++ b/src/usecases/timeline.js
@@ -52,22 +52,18 @@ module.exports.receiveMessage = async function(session, expectedMessage) {
         return getMessageFromEventTile(lastTile);
     }
 
-    let lastMessage = null;
-    let isExpectedMessage = false;
-    let totalTime = 0;
-    while (!isExpectedMessage) {
+    let lastMessage;
+    await session.poll(async () => {
         try {
             lastMessage = await getLastMessage();
-            isExpectedMessage = lastMessage &&
-                lastMessage.body === expectedMessage.body &&
-                lastMessage.sender === expectedMessage.sender
-        } catch(err) {}
-        if (totalTime > 5000) {
-            throw new Error("timed out after 5000ms");
+        } catch(err) {
+            return false;
         }
-        totalTime += 200;
-        await session.delay(200);
-    }
+        // stop polling when found the expected message
+        return lastMessage &&
+            lastMessage.body === expectedMessage.body &&
+            lastMessage.sender === expectedMessage.sender;
+    }, 5000, 200);
     assertMessage(lastMessage, expectedMessage);
     session.log.done();
 }