From e919335ed2216563d975dc8794eb863893b4daca Mon Sep 17 00:00:00 2001
From: Appurva Murawat <appurva21@proton.me>
Date: Fri, 19 May 2023 15:51:56 +0530
Subject: [PATCH] Add support for `pm.connection` and `pm.message`

---
 lib/sandbox/execute.js   | 13 ++++++++++++-
 lib/sandbox/execution.js |  3 ++-
 lib/sandbox/pmapi.js     | 19 +++++++++++++++++--
 3 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/lib/sandbox/execute.js b/lib/sandbox/execute.js
index e2321ff5..2c9907b3 100644
--- a/lib/sandbox/execute.js
+++ b/lib/sandbox/execute.js
@@ -138,6 +138,16 @@ module.exports = function (bridge, glob) {
 
                 bridge.dispatch(assertionEventName, options.cursor, assertions);
                 bridge.dispatch(EXECUTION_ASSERTION_EVENT, options.cursor, assertions);
+            },
+
+            // TODO: Handle timers
+            dispatchConnectionEvents = {
+                sendMessage (...args) {
+                    bridge.dispatch('connection.sendMessage.' + id, ...args);
+                },
+                close (...args) {
+                    bridge.dispatch('connection.close.' + id, ...args);
+                }
             };
 
         let waiting,
@@ -214,7 +224,8 @@ module.exports = function (bridge, glob) {
                     var eventId = timers.setEvent(callback);
 
                     bridge.dispatch(executionRequestEventName, options.cursor, id, eventId, request);
-                }, dispatchAssertions, new PostmanCookieStore(id, bridge, timers), {
+                }, dispatchAssertions, dispatchConnectionEvents,
+                new PostmanCookieStore(id, bridge, timers), {
                     disabledAPIs: initializationOptions.disabledAPIs
                 })
             ),
diff --git a/lib/sandbox/execution.js b/lib/sandbox/execution.js
index ab521ac9..ebe2dc1a 100644
--- a/lib/sandbox/execution.js
+++ b/lib/sandbox/execution.js
@@ -42,10 +42,11 @@ class Execution {
         });
 
         if (options.initializeExecution) {
-            const { request, response } = options.initializeExecution(this.target, context) || {};
+            const { request, response, message } = options.initializeExecution(this.target, context) || {};
 
             this.request = request;
             this.response = response;
+            this.message = message;
         }
         else {
             if (TARGETS_WITH_REQUEST[this.target] || _.has(context, PROPERTY.REQUEST)) {
diff --git a/lib/sandbox/pmapi.js b/lib/sandbox/pmapi.js
index 6e860f37..4a2d1910 100644
--- a/lib/sandbox/pmapi.js
+++ b/lib/sandbox/pmapi.js
@@ -45,11 +45,12 @@ const _ = require('lodash'),
  * @param {Execution} execution -
  * @param {Function} onRequest -
  * @param {Function} onAssertion -
+ * @param {Object} onConnection -
  * @param {Object} cookieStore -
  * @param {Object} [options] -
  * @param {Array.<String>} [options.disabledAPIs] -
  */
-function Postman (execution, onRequest, onAssertion, cookieStore, options = {}) {
+function Postman (execution, onRequest, onAssertion, onConnection, cookieStore, options = {}) {
     // @todo - ensure runtime passes data in a scope format
     let iterationData = new VariableScope();
 
@@ -167,6 +168,8 @@ function Postman (execution, onRequest, onAssertion, cookieStore, options = {})
          */
         response: execution.response,
 
+        message: execution.message,
+
         /**
          * The cookies object contains a list of cookies that are associated with the domain
          * to which the request was made.
@@ -250,7 +253,9 @@ function Postman (execution, onRequest, onAssertion, cookieStore, options = {})
             });
 
             return self;
-        }
+        },
+
+        connection: onConnection
     }, options.disabledAPIs);
 
     // extend pm api with test runner abilities
@@ -265,6 +270,16 @@ function Postman (execution, onRequest, onAssertion, cookieStore, options = {})
             }
         });
     }
+
+    if (this.message) {
+        // these are removed before serializing see `purse.js`
+        Object.defineProperty(this.message, 'to', {
+            get () {
+                return chai.expect(this).to;
+            }
+        });
+    }
+
     // add request assertions
     if (this.request) {
         // these are removed before serializing see `purse.js`