From ae6ec98e2ea8a4d7f167265a9115a44a35fb1753 Mon Sep 17 00:00:00 2001
From: Leonardo Venturini <leovenbag@gmail.com>
Date: Thu, 25 Jul 2024 18:03:42 -0400
Subject: [PATCH] make connection utility to capture messages

---
 packages/ddp-server/livedata_server_tests.js | 19 ++++--------
 packages/test-helpers/connection.js          | 31 ++++++++++++++++++++
 packages/test-helpers/package.js             | 30 ++++++++++++++-----
 3 files changed, 58 insertions(+), 22 deletions(-)

diff --git a/packages/ddp-server/livedata_server_tests.js b/packages/ddp-server/livedata_server_tests.js
index 57ed24acb4..3d79e64d6c 100644
--- a/packages/ddp-server/livedata_server_tests.js
+++ b/packages/ddp-server/livedata_server_tests.js
@@ -447,20 +447,9 @@ Tinytest.addAsync("livedata server - waiting for Promise", (test, onComplete) =>
  * https://github.com/meteor/meteor/issues/13212
  */
 Tinytest.addAsync('livedata server - publish cursor is properly awaited', async function (test) {
-  const messages = []
-
   let sub = null;
 
-  const { clientConn } = await getTestConnections(test)
-
-  const send = clientConn._stream.send
-
-  clientConn._stream.send = function (...args) {
-    send.apply(this, args)
-    messages.push(args[0])
-  }
-
-  clientConn._stream.on('message', message => messages.push(message));
+  const { conn, messages, cleanup } = await captureConnectionMessages(test);
 
   const coll = new Mongo.Collection('items', {
     defineMutationMethods: false,
@@ -482,7 +471,7 @@ Tinytest.addAsync('livedata server - publish cursor is properly awaited', async
   const reactiveVar = new ReactiveVar(1);
 
   const computation = Tracker.autorun(() => {
-    sub = clientConn.subscribe(publicationName, reactiveVar.get());
+    sub = conn.subscribe(publicationName, reactiveVar.get());
   });
 
   await sleep(100)
@@ -496,11 +485,13 @@ Tinytest.addAsync('livedata server - publish cursor is properly awaited', async
   /**
    * There shouldn't ever be `removed` messages here, otherwise the UI will glitch
    */
-  const parsedMessages = messages.map(m => EJSON.parse(m).msg)
+  const parsedMessages = messages.map(m => m.msg)
 
   test.equal(parsedMessages, expectedMessages)
 
   computation.stop();
+
+  cleanup()
 });
 
 function getTestConnections(test) {
diff --git a/packages/test-helpers/connection.js b/packages/test-helpers/connection.js
index a78c3dab46..4a360d1e97 100644
--- a/packages/test-helpers/connection.js
+++ b/packages/test-helpers/connection.js
@@ -52,3 +52,34 @@ makeTestConnection = function (test, succeeded, failed) {
     }
   );
 };
+
+createTestConnectionPromise = function (test) {
+  return new Promise((resolve, reject) => {
+    makeTestConnection(test, resolve, reject);
+  });
+};
+
+captureConnectionMessages = async function (test) {
+  const messages = []
+
+  const conn = await createTestConnectionPromise(test);
+
+  const send = conn._stream.send;
+
+  conn._stream.send = function (...args) {
+    send.apply(this, args);
+    messages.push(EJSON.parse(args[0]));
+  }
+
+  conn._stream.on('message', message => messages.push(EJSON.parse(message)));
+
+  function cleanup() {
+    conn._stream.send = send
+  }
+
+  return {
+    conn,
+    messages,
+    cleanup
+  }
+};
\ No newline at end of file
diff --git a/packages/test-helpers/package.js b/packages/test-helpers/package.js
index f03739d8ae..f50112d5c9 100644
--- a/packages/test-helpers/package.js
+++ b/packages/test-helpers/package.js
@@ -22,15 +22,29 @@ Package.onUse(function (api) {
   // the like.
   api.use('ddp');
 
-
   api.export([
-    'pollUntil', 'try_all_permutations',
-    'SeededRandom', 'clickElement', 'blurElement',
-    'focusElement', 'simulateEvent', 'getStyleProperty', 'canonicalizeHtml',
-    'renderToDiv', 'clickIt',
-    'withCallbackLogger', 'testAsyncMulti',
-    'simplePoll', 'runAndThrowIfNeeded',
-    'makeTestConnection', 'DomUtils', 'mockBehaviours', 'waitUntil']);
+    'pollUntil',
+    'try_all_permutations',
+    'SeededRandom',
+    'clickElement',
+    'blurElement',
+    'focusElement',
+    'simulateEvent',
+    'getStyleProperty',
+    'canonicalizeHtml',
+    'renderToDiv',
+    'clickIt',
+    'withCallbackLogger',
+    'testAsyncMulti',
+    'simplePoll',
+    'runAndThrowIfNeeded',
+    'DomUtils',
+    'mockBehaviours',
+    'waitUntil',
+    'makeTestConnection',
+    'createTestConnectionPromise',
+    'captureConnectionMessages',
+  ]);
 
   api.addFiles('try_all_permutations.js');
   api.addFiles('async_multi.js');