From 491f5b0685136602cb204f8e152faaeef503f163 Mon Sep 17 00:00:00 2001
From: Anna Henningsen <anna.henningsen@mongodb.com>
Date: Thu, 7 Mar 2024 14:15:55 +0100
Subject: [PATCH] fix(cli-repl): properly report API telemetry for plain-vm
 eval (#1859)

---
 packages/cli-repl/src/async-repl.ts    |  2 +-
 packages/cli-repl/src/cli-repl.spec.ts | 37 ++++++++++++++++++++++++--
 packages/cli-repl/src/mongosh-repl.ts  |  4 +++
 3 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/packages/cli-repl/src/async-repl.ts b/packages/cli-repl/src/async-repl.ts
index 78fd953cd..c05ee3d7c 100644
--- a/packages/cli-repl/src/async-repl.ts
+++ b/packages/cli-repl/src/async-repl.ts
@@ -81,7 +81,7 @@ export function start(opts: AsyncREPLOptions): REPLServer {
     (opts as ReplOptions).breakEvalOnSigint = true;
   }
 
-  const repl = (opts.start ?? originalStart)(opts);
+  const repl: REPLServer = (opts.start ?? originalStart)(opts);
   const originalEval = promisify(
     wrapPauseInput(repl.input, wrapNoSyncDomainError(repl.eval.bind(repl)))
     // string, Context, string is not string, any, string
diff --git a/packages/cli-repl/src/cli-repl.spec.ts b/packages/cli-repl/src/cli-repl.spec.ts
index af2235202..8ed447eeb 100644
--- a/packages/cli-repl/src/cli-repl.spec.ts
+++ b/packages/cli-repl/src/cli-repl.spec.ts
@@ -1288,15 +1288,16 @@ describe('CliRepl', function () {
           setTelemetryDelay(0);
         });
 
-        it('timeouts fast', async function () {
+        it('times out fast', async function () {
+          const testStartMs = Date.now();
           // The `httpRequestTimeout` of our Segment Analytics is set to
           // 1000ms which makes these requests fail as they exceed the timeout.
           // Segment will silently fail http request errors when tracking and flushing.
           // This will cause the test to fail if the system running the tests is
           //  unable to flush telemetry in 1500ms.
-          this.timeout(2500);
           setTelemetryDelay(5000);
           await cliRepl.start(await testServer.connectionString(), {});
+          this.timeout(Date.now() - testStartMs + 2500); // Do not include connection time in 2.5s timeout
           input.write('use somedb;\n');
           input.write('exit\n');
           await waitBus(cliRepl.bus, 'mongosh:closed');
@@ -1451,6 +1452,38 @@ describe('CliRepl', function () {
           expect(totalEventsTracked).to.equal(5);
         });
 
+        it('sends out telemetry data for multiple command line scripts', async function () {
+          cliReplOptions.shellCliOptions.eval = [
+            'db.hello(); db.hello();',
+            'db.hello()',
+          ];
+          cliRepl = new CliRepl(cliReplOptions);
+          await startWithExpectedImmediateExit(
+            cliRepl,
+            await testServer.connectionString()
+          );
+          expect(totalEventsTracked).to.equal(7);
+
+          const apiEvents = requests
+            .map((req) =>
+              JSON.parse(req.body).batch.filter(
+                (entry) => entry.event === 'API Call'
+              )
+            )
+            .flat();
+          expect(apiEvents).to.have.lengthOf(2);
+          expect(
+            apiEvents.map((e) => [
+              e.properties.class,
+              e.properties.method,
+              e.properties.count,
+            ])
+          ).to.deep.equal([
+            ['Database', 'hello', 2],
+            ['Database', 'hello', 1],
+          ]);
+        });
+
         it('sends out telemetry if the repl is running in an interactive mode in a containerized environment', async function () {
           cliRepl = new CliRepl(cliReplOptions);
           cliRepl.getIsContainerizedEnvironment = () => {
diff --git a/packages/cli-repl/src/mongosh-repl.ts b/packages/cli-repl/src/mongosh-repl.ts
index b6e16618f..b4067a859 100644
--- a/packages/cli-repl/src/mongosh-repl.ts
+++ b/packages/cli-repl/src/mongosh-repl.ts
@@ -503,6 +503,8 @@ class MongoshNodeRepl implements EvaluationListener {
     }
 
     markTime(TimingCategories.UserConfigLoading, 'set up history file');
+
+    // Similar calls are added for plain-vm evaluation below
     (repl as any).on(asyncRepl.evalStart, () => {
       this.bus.emit('mongosh:evaluate-started');
     });
@@ -746,6 +748,7 @@ class MongoshNodeRepl implements EvaluationListener {
         );
       };
     });
+    this.bus.emit('mongosh:evaluate-started');
     try {
       process.addListener('SIGINT', asyncSigintHandler);
       return await Promise.race([
@@ -762,6 +765,7 @@ class MongoshNodeRepl implements EvaluationListener {
       ]);
     } finally {
       process.removeListener('SIGINT', asyncSigintHandler);
+      this.bus.emit('mongosh:evaluate-finished');
     }
   }