fix: custom states will delay if added in constructor

Author: calebdwilliams

Diff for: src/CustomStateSet.ts

@@ -3,6 +3,13 @@ import { ICustomElement } from "./types";
 /** Save a reference to the ref for teh CustomStateSet */
 const customStateMap = new WeakMap<CustomStateSet, ICustomElement>();
 
+function addState(ref: ICustomElement, stateName: string): void {
+  ref.toggleAttribute(stateName, true);
+  if (ref.part) {
+    ref.part.add(stateName);
+  }
+}
+
 export type CustomState = `--${string}`;
 
 export class CustomStateSet extends Set<CustomState> {
@@ -25,10 +32,21 @@ export class CustomStateSet extends Set<CustomState> {
     }
     const result = super.add(state);
     const ref = customStateMap.get(this);
-    ref.toggleAttribute(`state${state}`, true);
-    if (ref.part) {
-      ref.part.add(`state${state}`);
+    const stateName = `state${state}`;
+
+    /**
+     * Only add the state immediately if the ref is connected to the DOM;
+     * otherwise, wait a tick because the element is likely being constructed
+     * by document.createElement and would throw otherwise.
+     */
+    if (ref.isConnected) {
+      addState(ref, stateName);
+    } else {
+      setTimeout(() => {
+        addState(ref, stateName);
+      });
     }
+
     return result;
   }
 

Diff for: test/CustomStateSet.test.ts

@@ -1,7 +1,21 @@
-import { fixture, html, expect, fixtureCleanup } from '@open-wc/testing';
+import '../dist/index.js';
+
+import { fixture, html, expect, fixtureCleanup, aTimeout } from '@open-wc/testing';
 import { ICustomElement } from '../dist';
 import { CustomStateSet } from '../src/CustomStateSet';
 
+class CustomElementStateInConstructor extends HTMLElement {
+  internals = this.attachInternals();
+  constructor() {
+    super();
+
+    this.internals.states.add('--foo');
+  }
+}
+
+const tagName = 'custom-element-state-in-constructor'
+customElements.define(tagName, CustomElementStateInConstructor);
+
 describe('CustomStateSet polyfill', () => {
   let el: HTMLElement;
   let set: CustomStateSet;
@@ -56,4 +70,18 @@ describe('CustomStateSet polyfill', () => {
       expect(el.part.contains('state--bar')).to.be.false;
     }
   });
+
+  it('will use a timeout if a state is added in a constructor', async () => {
+    let el;
+    expect(() => {
+      el = document.createElement(tagName);
+    }).not.to.throw();
+
+    if (window.CustomStateSet.isPolyfilled) {
+      await aTimeout(100);
+      expect(el.matches('[state--foo]')).to.be.true;
+    } else {
+      expect(el.matches(`:--foo`)).to.be.true;
+    }
+  });
 });