Fix the issues with post-download upgrade (#3192)

src/custom-element.js

@@ -81,6 +81,7 @@ export function upgradeOrRegisterElement(win, name, toClass) {
   user.assert(knownElements[name] == ElementStub,
       '%s is already registered. The script tag for ' +
       '%s is likely included twice in the page.', name, name);
+  knownElements[name] = toClass;
   for (let i = 0; i < stubbedElements.length; i++) {
     const stub = stubbedElements[i];
     // There are 3 possible states here:
@@ -279,10 +280,11 @@ class AmpElement {
  *
  * @param {!Window} win The window in which to register the elements.
  * @param {string} name Name of the custom element
- * @param {function(new:BaseElement, !Element)} implementationClass
+ * @param {function(new:BaseElement, !Element)} opt_implementationClass For
+ *     testing only.
  * @return {!AmpElement.prototype}
  */
-export function createAmpElementProto(win, name, implementationClass) {
+export function createAmpElementProto(win, name, opt_implementationClass) {
   /**
    * @lends {AmpElement.prototype}
    */
@@ -352,8 +354,11 @@ export function createAmpElementProto(win, name, implementationClass) {
     /** @private {?Element|undefined} */
     this.overflowElement_ = undefined;
 
+    // `opt_implementationClass` is only used for tests.
+    const Ctor = opt_implementationClass || knownElements[name];
+
     /** @private {!BaseElement} */
-    this.implementation_ = new implementationClass(this);
+    this.implementation_ = new Ctor(this);
     this.implementation_.createdCallback();
 
     /**
@@ -1168,7 +1173,7 @@ export function registerElement(win, name, implementationClass) {
   knownElements[name] = implementationClass;
 
   win.document.registerElement(name, {
-    prototype: createAmpElementProto(win, name, implementationClass),
+    prototype: createAmpElementProto(win, name),
   });
 }
 
@@ -1184,8 +1189,9 @@ export function registerElementAlias(win, aliasName, sourceName) {
   const implementationClass = knownElements[sourceName];
 
   if (implementationClass) {
+    knownElements[aliasName] = implementationClass;
     win.document.registerElement(aliasName, {
-      prototype: createAmpElementProto(win, aliasName, implementationClass),
+      prototype: createAmpElementProto(win, aliasName),
     });
   } else {
     throw new Error(`Element name is unknown: ${sourceName}.` +
@@ -1199,6 +1205,7 @@ export function registerElementAlias(win, aliasName, sourceName) {
  * This makes it possible to mark an element as loaded in a test.
  * @param {!Window} win
  * @param {string} elementName Name of an extended custom element.
+ * @visibleForTesting
  */
 export function markElementScheduledForTesting(win, elementName) {
   if (!win.ampExtendedElements) {
@@ -1211,6 +1218,7 @@ export function markElementScheduledForTesting(win, elementName) {
  * Resets our scheduled elements.
  * @param {!Window} win
  * @param {string} elementName Name of an extended custom element.
+ * @visibleForTesting
  */
 export function resetScheduledElementForTesting(win, elementName) {
   if (win.ampExtendedElements) {
@@ -1219,3 +1227,12 @@ export function resetScheduledElementForTesting(win, elementName) {
   delete knownElements[elementName];
 }
 
+/**
+ * Returns a currently registered element class.
+ * @param {string} elementName Name of an extended custom element.
+ * @return {?function()}
+ * @visibleForTesting
+ */
+export function getElementClassForTesting(elementName) {
+  return knownElements[elementName] || null;
+}

test/functional/test-custom-element.js

@@ -24,9 +24,12 @@ import * as sinon from 'sinon';
 import {getService, resetServiceForTesting} from '../../src/service';
 import {
   createAmpElementProto,
+  getElementClassForTesting,
   markElementScheduledForTesting,
+  registerElement,
   resetScheduledElementForTesting,
   stubElements,
+  upgradeOrRegisterElement,
 } from '../../src/custom-element';
 // TODO(@cramforce): Move tests into their own file.
 import {
@@ -35,6 +38,87 @@ import {
 } from '../../src/element-service';
 
 
+describe('CustomElement register', () => {
+
+  class ConcreteElement extends BaseElement {}
+
+  let sandbox;
+  let win;
+  let doc;
+  let registeredElements = {};
+
+  beforeEach(() => {
+    sandbox = sinon.sandbox.create();
+    resetScheduledElementForTesting(window, 'amp-element1');
+
+    win = {
+      Object: window.Object,
+      HTMLElement: window.HTMLElement,
+      services: window.services,
+    };
+
+    registeredElements = {};
+    doc = {
+      registerElement: (name, spec) => {
+        if (registeredElements[name]) {
+          throw new Error('already registered: ' + name);
+        }
+        registeredElements[name] = spec;
+      },
+    };
+    win.document = doc;
+  });
+
+  afterEach(() => {
+    sandbox.restore();
+    resetScheduledElementForTesting(window, 'amp-element1');
+  });
+
+  function createElement(elementName) {
+    const spec = registeredElements[elementName];
+    if (!spec) {
+      throw new Error('unknown element: ' + elementName);
+    }
+    let ctor = spec.ctor;
+    if (!ctor) {
+      const proto = spec.prototype;
+      ctor = function() {
+        const el = document.createElement(elementName);
+        el.__proto__ = proto;
+        return el;
+      };
+      ctor.prototype = proto;
+      proto.constructor = ctor;
+      spec.ctor = ctor;
+    }
+    const element = new ctor();
+    element.createdCallback();
+    return element;
+  }
+
+  it('should go through stub/upgrade cycle', () => {
+    registerElement(win, 'amp-element1', ElementStub);
+    expect(getElementClassForTesting('amp-element1')).to.equal(ElementStub);
+    expect(registeredElements['amp-element1']).to.exist;
+    expect(registeredElements['amp-element1'].prototype).to.exist;
+
+    // Pre-download elements are created as ElementStub.
+    const element1 = createElement('amp-element1');
+    expect(element1.implementation_).to.be.instanceOf(ElementStub);
+
+    // Post-download, elements are upgraded.
+    upgradeOrRegisterElement(win, 'amp-element1', ConcreteElement);
+    expect(getElementClassForTesting('amp-element1')).to.equal(ConcreteElement);
+    expect(element1.implementation_).to.be.instanceOf(ConcreteElement);
+
+    // Elements created post-download and immediately upgraded.
+    const element2 = createElement('amp-element1');
+    element2.createdCallback();
+    expect(element2.implementation_).to.be.instanceOf(ConcreteElement);
+  });
+});
+
+
 describe('CustomElement', () => {
 
   const resources = resourcesFor(window);