From 2cb04f0c22d00ed22411a4543e03714096a24578 Mon Sep 17 00:00:00 2001
From: Micajuine Ho <micajuineho@google.com>
Date: Thu, 3 Sep 2020 11:54:34 -0700
Subject: [PATCH] Clean and add error msg

---
 extensions/amp-analytics/0.1/events.js        | 32 ++++++--------
 .../amp-analytics/0.1/test/test-events.js     | 43 ++++++++++++-------
 2 files changed, 40 insertions(+), 35 deletions(-)

diff --git a/extensions/amp-analytics/0.1/events.js b/extensions/amp-analytics/0.1/events.js
index 100d5469fac24..598068344062e 100644
--- a/extensions/amp-analytics/0.1/events.js
+++ b/extensions/amp-analytics/0.1/events.js
@@ -1489,7 +1489,7 @@ export class VisibilityTracker extends EventTracker {
       // When `selector` is specified, we always use "ini-load" signal as
       // a "ready" signal.
       const readyPromiseWaitForSpec =
-        waitForSpec || (selector ? 'ini-load' : null);
+        waitForSpec || (selector ? 'ini-load' : 'none');
       return visibilityManager.listenRoot(
         visibilitySpec,
         this.getReadyPromise(readyPromiseWaitForSpec),
@@ -1633,20 +1633,6 @@ export class VisibilityTracker extends EventTracker {
     return 'onpagehide' in this.root.ampdoc.win;
   }
 
-  /**
-   * Sets the default waitForSpec based upon if the element
-   * passed in is an AMP element or not.
-   * @param {string|undefined} waitForSpec
-   * @param {!Element} element
-   * @return {string}
-   */
-  setDefaultWaitForElement_(waitForSpec, element) {
-    if (!waitForSpec) {
-      waitForSpec = isAmpElement(element) ? 'ini-load' : 'none';
-    }
-    return waitForSpec;
-  }
-
   /**
    * @param {string|undefined} waitForSpec
    * @param {Element=} opt_element
@@ -1655,19 +1641,25 @@ export class VisibilityTracker extends EventTracker {
    */
   getReadyPromise(waitForSpec, opt_element) {
     if (opt_element) {
-      waitForSpec = this.setDefaultWaitForElement_(waitForSpec, opt_element);
+      if (!isAmpElement(opt_element)) {
+        userAssert(
+          !waitForSpec || waitForSpec == 'none',
+          'waitFor for non-AMP elements must be none or null. Found %s',
+          waitForSpec
+        );
+      } else {
+        waitForSpec = waitForSpec || 'ini-load';
+      }
     }
 
-    if (!waitForSpec) {
+    if (!waitForSpec || waitForSpec == 'none') {
       // Default case, waitFor selector is not defined, wait for nothing
       return null;
     }
 
     const trackerAllowlist = getTrackerTypesForParentType('visible');
-    const isAllowedWaitFor = opt_element ? isAmpElement(opt_element) : true;
     userAssert(
-      waitForSpec == 'none' ||
-        (trackerAllowlist[waitForSpec] !== undefined && isAllowedWaitFor),
+      trackerAllowlist[waitForSpec] !== undefined,
       'waitFor value %s not supported',
       waitForSpec
     );
diff --git a/extensions/amp-analytics/0.1/test/test-events.js b/extensions/amp-analytics/0.1/test/test-events.js
index a661844bec28d..27f130156e69c 100644
--- a/extensions/amp-analytics/0.1/test/test-events.js
+++ b/extensions/amp-analytics/0.1/test/test-events.js
@@ -2247,32 +2247,45 @@ describes.realWin('Events', {amp: 1}, (env) => {
       });
 
       describe('non AMP elements', () => {
-        it('with non AMP element and waitFor NONE', () => {
+        it('with non AMP element and waitFor NONE or null', () => {
           const element = win.document.createElement('p');
           expect(tracker.getReadyPromise('none', element)).to.be.null;
+          expect(tracker.getReadyPromise(null, element)).to.be.null;
         });
 
-        it('error with non AMP element and waitFor not NONE', () => {
+        it('error with non AMP element and waitFor not NONE or null', () => {
           const element = win.document.createElement('p');
           expect(() => tracker.getReadyPromise('ini-load', element)).to.throw(
-            /waitFor value ini-load not supported​​​/
-          );
-        });
-      });
-
-      describe('default waitFor with element', () => {
-        it('should set default waitFor for non AMP element', () => {
-          const element = win.document.createElement('p');
-          expect(tracker.setDefaultWaitForElement_(null, element)).to.equal(
-            'none'
+            /waitFor for non-AMP elements must be none or null. Found ini-load/
           );
         });
 
-        it('should set default waitFor for AMP element', () => {
+        it('should set default waitFor for AMP element', async () => {
           const element = win.document.createElement('amp-list');
-          expect(tracker.setDefaultWaitForElement_(null, element)).to.equal(
-            'ini-load'
+          tracker.waitForTrackers_['render-start'] = root.getTracker(
+            'render-start',
+            SignalTracker
           );
+          const signalTrackerMock = env.sandbox.mock(
+            tracker.waitForTrackers_['render-start']
+          );
+          const iniLoadTrackerMock = env.sandbox.mock(
+            tracker.waitForTrackers_['ini-load']
+          );
+
+          await tracker.getReadyPromise(null, element);
+          iniLoadTrackerMock
+            .expects('getElementSignal')
+            .withExactArgs('ini-load', element)
+            .returns(Promise.resolve())
+            .once();
+
+          await tracker.getReadyPromise('render-start', element);
+          signalTrackerMock
+            .expects('getElementSignal')
+            .withExactArgs('render-start')
+            .returns(Promise.resolve())
+            .once();
         });
       });
     });