diff --git a/package.json b/package.json
index 7329988..e61a17e 100644
--- a/package.json
+++ b/package.json
@@ -9,7 +9,8 @@
     "lint": "npm run lint:js",
     "dev": "wrangler dev src/edge.js",
     "deploy": "wrangler deploy src/edge.js",
-    "test": "c8 mocha"
+
+    "test": "c8 mocha --exit # --exit is needed because some test code triggers async listeners"
   },
   "mocha": {
     "reporter": "mocha-multi-reporters",
diff --git a/src/edge.js b/src/edge.js
index 374bc8f..be68597 100644
--- a/src/edge.js
+++ b/src/edge.js
@@ -215,7 +215,7 @@ export const updateHandler = (update, _origin, doc) => {
   doc.conns.forEach((_, conn) => send(doc, conn, message));
 };
 
-class WSSharedDoc extends Y.Doc {
+export class WSSharedDoc extends Y.Doc {
   constructor(name) {
     super({ gc: gcEnabled });
     this.name = name;
diff --git a/test/edge.test.js b/test/edge.test.js
index 95d8f85..6bfc93e 100644
--- a/test/edge.test.js
+++ b/test/edge.test.js
@@ -10,7 +10,32 @@
  * governing permissions and limitations under the License.
  */
 import assert from 'assert';
-import { updateHandler } from '../src/edge.js';
+import { updateHandler, WSSharedDoc } from '../src/edge.js';
+
+function isSubArray(full, sub) {
+  if (sub.length === 0) {
+    return true;
+  }
+
+  const candidateIdxs = [];
+  for (let i = 0; i < full.length; i++) {
+    if (full[i] === sub[0]) {
+      candidateIdxs.push(i);
+    }
+  }
+
+  nextCandidate:
+  for (let i = 0; i < candidateIdxs.length; i++) {
+    for (let j = 0; j < sub.length; j++) {
+      if (sub[j] !== full[candidateIdxs[i] + j]) {
+        break nextCandidate;
+      }
+    }
+    return true;
+  }
+
+  return false;
+}
 
 describe('Collab Test Suite', () => {
   it('Test updateHandler', () => {
@@ -29,13 +54,11 @@ describe('Collab Test Suite', () => {
       },
     };
 
-    const fe = (func) => {
-      func(null, conn);
-    };
-
     const deleted = [];
     const conns = {
-      forEach: fe,
+      forEach(f) {
+        f(null, conn);
+      },
       has(c) {
         return c === conn;
       },
@@ -97,4 +120,31 @@ describe('Collab Test Suite', () => {
     assert.deepStrictEqual(update, conn1.message.slice(-4));
     assert.deepStrictEqual(update, conn2.message.slice(-4));
   });
+
+  it('Test WSSharedDoc', () => {
+    const doc = new WSSharedDoc('hello');
+    assert.equal(doc.name, 'hello');
+    assert.equal(doc.awareness.getLocalState(), null);
+
+    const conn = {
+      isClosed: false,
+      message: null,
+      readyState: 1, // wsReadyStateOpen
+      has() {
+        return true;
+      },
+      close() {
+        this.isClosed = true;
+      },
+      send(m) {
+        this.message = m;
+      },
+    };
+
+    doc.conns.set(conn, 'conn1');
+    doc.awareness.setLocalState('foo');
+    assert(conn.isClosed === false);
+    const fooAsUint8Arr = new Uint8Array(['f'.charCodeAt(0), 'o'.charCodeAt(0), 'o'.charCodeAt(0)]);
+    assert(isSubArray(conn.message, fooAsUint8Arr));
+  });
 });