Await overriden functions when returning promises

addon/-private/apollo/setup-hooks.js

@@ -27,8 +27,15 @@ function installHook(queryManager, context, hookName) {
 
   context[hookName] = function () {
     if (typeof originalHook === 'function') {
-      originalHook.call(this, ...arguments);
+      const result = originalHook.call(this, ...arguments);
+
+      if (result instanceof Promise) {
+        return result.then(() => {
+          hook.call(queryManager, ...arguments);
+        });
+      }
     }
+
     hook.call(queryManager, ...arguments);
   };
 }

tests/unit/-private/query-manager-hooks-route-test.js

@@ -39,85 +39,120 @@ module('Unit | queryManager | Setup Hooks in route', function (hooks) {
   hooks.beforeEach(function (assert) {
     assertDeepEqual = assert.deepEqual.bind(assert);
     unsubscribeCalled = 0;
-
-    this.subject = function () {
-      this.owner.register('service:overridden-apollo', OverriddenApollo);
-      this.owner.register('test-container:test-object', TestObject);
-      return this.owner.lookup('test-container:test-object');
-    };
   });
 
-  test('it unsubscribes from any watchQuery subscriptions with isExiting=true', function (assert) {
-    let subject = this.subject();
+  module('with synchronous overridden functions', function (innerHooks) {
+    innerHooks.beforeEach(function () {
+      this.subject = function () {
+        this.owner.register('service:overridden-apollo', OverriddenApollo);
+        this.owner.register('test-container:test-object', TestObject);
+        return this.owner.lookup('test-container:test-object');
+      };
+    });
 
-    assert.equal(unsubscribeCalled, 0, 'should have been initialized with 0');
-    subject.apollo.watchQuery({ query: 'fakeQuery' });
-    subject.apollo.watchQuery({ query: 'fakeQuery' });
+    test('it unsubscribes from any watchQuery subscriptions with isExiting=true', function (assert) {
+      let subject = this.subject();
 
-    subject.beforeModel();
-    subject.resetController({}, true);
-    assert.equal(
-      unsubscribeCalled,
-      2,
-      '_apolloUnsubscribe() was called once per watchQuery'
-    );
-  });
+      assert.equal(unsubscribeCalled, 0, 'should have been initialized with 0');
+      subject.apollo.watchQuery({ query: 'fakeQuery' });
+      subject.apollo.watchQuery({ query: 'fakeQuery' });
 
-  test('it unsubscribes from any subscriptions', function (assert) {
-    let subject = this.subject();
+      subject.beforeModel();
+      subject.resetController({}, true);
+      assert.equal(
+        unsubscribeCalled,
+        2,
+        '_apolloUnsubscribe() was called once per watchQuery'
+      );
+    });
 
-    assert.equal(unsubscribeCalled, 0, 'should have been initialized with 0');
-    subject.apollo.subscribe({ query: 'fakeSubscription' });
-    subject.apollo.subscribe({ query: 'fakeSubscription' });
+    test('it unsubscribes from any subscriptions', function (assert) {
+      let subject = this.subject();
 
-    subject.beforeModel();
-    subject.resetController({}, true);
-    assert.equal(
-      unsubscribeCalled,
-      2,
-      '_apolloUnsubscribe() was called once per subscribe'
-    );
-  });
+      assert.equal(unsubscribeCalled, 0, 'should have been initialized with 0');
+      subject.apollo.subscribe({ query: 'fakeSubscription' });
+      subject.apollo.subscribe({ query: 'fakeSubscription' });
 
-  test('it only unsubscribes from stale watchQuery subscriptions with isExiting=false', function (assert) {
-    let subject = this.subject();
-
-    assert.equal(unsubscribeCalled, 0, 'should have been initialized with 0');
-    subject.apollo.watchQuery({ query: 'fakeQuery' });
-
-    // simulate data being re-fetched, as when query params change
-    subject.beforeModel();
-    subject.apollo.watchQuery({ query: 'fakeQuery' });
-
-    subject.resetController({}, false);
-    assert.equal(
-      unsubscribeCalled,
-      1,
-      '_apolloUnsubscribe() was called only once, for the first query'
-    );
-
-    subject.beforeModel();
-    subject.willDestroy();
-    assert.equal(
-      unsubscribeCalled,
-      2,
-      '_apolloUnsubscribe() was called for all quries'
-    );
-  });
+      subject.beforeModel();
+      subject.resetController({}, true);
+      assert.equal(
+        unsubscribeCalled,
+        2,
+        '_apolloUnsubscribe() was called once per subscribe'
+      );
+    });
 
-  test('it unsubscribes from any watchQuery subscriptions on willDestroy', function (assert) {
-    let subject = this.subject();
+    test('it only unsubscribes from stale watchQuery subscriptions with isExiting=false', function (assert) {
+      let subject = this.subject();
+
+      assert.equal(unsubscribeCalled, 0, 'should have been initialized with 0');
+      subject.apollo.watchQuery({ query: 'fakeQuery' });
+
+      // simulate data being re-fetched, as when query params change
+      subject.beforeModel();
+      subject.apollo.watchQuery({ query: 'fakeQuery' });
+
+      subject.resetController({}, false);
+      assert.equal(
+        unsubscribeCalled,
+        1,
+        '_apolloUnsubscribe() was called only once, for the first query'
+      );
+
+      subject.beforeModel();
+      subject.willDestroy();
+      assert.equal(
+        unsubscribeCalled,
+        2,
+        '_apolloUnsubscribe() was called for all quries'
+      );
+    });
 
-    assert.equal(unsubscribeCalled, 0, 'should have been initialized with 0');
-    subject.apollo.watchQuery({ query: 'fakeQuery' });
-    subject.apollo.watchQuery({ query: 'fakeQuery' });
+    test('it unsubscribes from any watchQuery subscriptions on willDestroy', function (assert) {
+      let subject = this.subject();
 
-    subject.beforeModel();
-    subject.willDestroy();
-    assert.equal(
-      unsubscribeCalled,
-      2,
-      '_apolloUnsubscribe() was called once per watchQuery'
-    );
+      assert.equal(unsubscribeCalled, 0, 'should have been initialized with 0');
+      subject.apollo.watchQuery({ query: 'fakeQuery' });
+      subject.apollo.watchQuery({ query: 'fakeQuery' });
+
+      subject.beforeModel();
+      subject.willDestroy();
+      assert.equal(
+        unsubscribeCalled,
+        2,
+        '_apolloUnsubscribe() was called once per watchQuery'
+      );
+    });
+  });
+
+  module('with asynchronous beforeModel', function (innerHooks) {
+    innerHooks.beforeEach(function () {
+      this.subject = function () {
+        class TestAsyncObject extends Route {
+          afterAsync = false;
+          @queryManager({ service: 'overridden-apollo' }) apollo;
+          async beforeModel() {
+            await new Promise((resolve) => setTimeout(resolve, 10));
+            this.afterAsync = true;
+            super.beforeModel();
+          }
+        }
+
+        this.owner.register('service:overridden-apollo', OverriddenApollo);
+        this.owner.register('test-container:test-object', TestAsyncObject);
+        return this.owner.lookup('test-container:test-object');
+      };
+    });
+
+    test('it must allow to await for the beforeModel function', async function (assert) {
+      assert.equal(unsubscribeCalled, 0, 'should have been initialized with 0');
+      const subject = this.subject();
+      subject.apollo.watchQuery({ query: 'fakeQuery' });
+      await subject.beforeModel();
+      assert.ok(
+        subject.afterAsync,
+        'original implementation of beforeModel should be awaited'
+      );
+    });
   });
 });