chore(e2e-tests): force wait for initialization before writing input (#2600)

In the past, we've had a number of flaky tests that followed a pattern
similar to this:

1. Start a test shell
2. Run a command using `executeLine()`
3. Make assertions about the command's output

The internal logic of `executeLine()` determines that it has seen
the end of the command's output by looking for the subsequent prompt
(similar to how a human being in front of the shell would determine
that the command has finished).

However, in the above scheme, there is a race condition that can
develop as follows:

1. The test shell process starts and prints part of its initial output
2. The `executeLine()` command remembers the current output position
   and sends the input to the process's stdin stream
3. The test shell prints its first prompt
4. The test reads the output from the shell and sees the prompt,
   assumes that the comamnd has finished
5. The test shell finishes evaluating the input and writes the
   corresponding output + a new prompt
6. The test ignores the remaining output and instead only treats
   the text up to the first prompt as the command output
7. An assertion made about the output fails, as the output is not
   coming from the command it was supposed to come from

In order to avoid the likelihood of introducing such race conditions,
this commit adds a guardrail to the input writing methods that
ensures that the shell has reached some form of known state that has
been awaited, such as seeing the first prompt or other specific output,
before it allows writing to the process's input.

Author: addaleax

packages/e2e-tests/test/e2e-direct.spec.ts

@@ -86,8 +86,10 @@ describe('e2e direct connection', function () {
         const shell = this.startTestShell({
           args: [await rs0.connectionString()],
         });
+        await shell.waitForPrompt();
         await shell.executeLine(`db.getSiblingDB("${dbname}").dropDatabase()`);
         shell.writeInputLine('exit');
+        await shell.waitForSuccessfulExit();
       });
 
       context('connecting to secondary members directly', function () {

packages/e2e-tests/test/e2e-fle.spec.ts

@@ -135,6 +135,7 @@ describe('FLE tests', function () {
           },
           cwd: path.join(__dirname, '..', '..', 'cli-repl', 'test', 'fixtures'),
         });
+        await shell.waitForPrompt();
 
         if (withEnvVarCredentials) {
           // Need to set up the AWS context inside the shell for enabling
@@ -926,6 +927,7 @@ describe('FLE tests', function () {
             await testServer.connectionString(),
           ],
         });
+        await shell.waitForPrompt();
         await shell.executeLine(`{
       const keyMongo = Mongo(
         db.getMongo(),

packages/e2e-tests/test/e2e.spec.ts

@@ -310,10 +310,9 @@ describe('e2e', function () {
       // The number of newlines here matters
       shell.writeInput(
         'sleep(100);print([1,2,3,4,5,6,7,8,9,10].reduce(\n(a,b) => { return a*b; }, 1))\n\n\n\n',
-        { end: true }
+        { end: true, requireFinishedInitialization: false }
       );
-      await shell.waitForSuccessfulExit();
-      shell.assertContainsOutput('3628800');
+      expect(await shell.waitForCleanOutput()).to.include('3628800');
     });
     it('ignores control characters in TTY input', async function () {
       shell = this.startTestShell({
@@ -1155,10 +1154,9 @@ describe('e2e', function () {
       });
       shell.writeInput(
         '[db.hello()].reduce(\n() => { return 11111*11111; },0)\n\n\n',
-        { end: true }
+        { end: true, requireFinishedInitialization: false }
       );
-      await shell.waitForSuccessfulExit();
-      shell.assertContainsOutput('123454321');
+      expect(await shell.waitForCleanOutput()).to.include('123454321');
     });
   });
 
@@ -2605,9 +2603,7 @@ describe('e2e', function () {
 
     it('allows connecting to a host and running commands against it', async function () {
       const connectionString = await testServer.connectionString();
-      await eventually(() => {
-        shell.assertContainsOutput('Please enter a MongoDB connection string');
-      });
+      await shell.waitForLine(/Please enter a MongoDB connection string/);
       shell.writeInputLine(connectionString);
       await shell.waitForPrompt();
 

packages/e2e-tests/test/test-shell.ts

@@ -20,6 +20,16 @@ type SignalType = ChildProcess extends { kill: (signal: infer T) => any }
   ? T
   : never;
 
+export interface TestShellInputOptions {
+  end?: boolean;
+  requireFinishedInitialization?: boolean;
+}
+
+export interface TestShellWaitForPromptOptions {
+  timeout?: number;
+  promptPattern?: RegExp;
+}
+
 // Assume that prompt strings are those that end in '> ' but do not contain
 // < or > (so that e.g. '- <repl>' in a stack trace is not considered a prompt).
 const PROMPT_PATTERN = /^([^<>]*> ?)+$/m;
@@ -152,6 +162,7 @@ export class TestShell {
   private _output: string;
   private _rawOutput: string;
   private _onClose: Promise<number>;
+  private _initializationKnownToBeFinished = false;
 
   constructor(
     shellProcess: ChildProcessWithoutNullStreams,
@@ -205,11 +216,14 @@ export class TestShell {
         });
       }
     });
+    // Not technically true, but in practice the patterns we wait for
+    // are sufficient to indicate that initialization is done.
+    this._initializationKnownToBeFinished = true;
   }
 
   async waitForPrompt(
     start = 0,
-    opts: { timeout?: number; promptPattern?: RegExp } = {}
+    opts: TestShellWaitForPromptOptions = {}
   ): Promise<void> {
     await eventually(
       () => {
@@ -247,10 +261,13 @@ export class TestShell {
       },
       { ...opts }
     );
+    this._initializationKnownToBeFinished = true;
   }
 
-  waitForAnyExit(): Promise<number> {
-    return this._onClose;
+  async waitForAnyExit(): Promise<number> {
+    const code = await this._onClose;
+    this._initializationKnownToBeFinished = true;
+    return code;
   }
 
   async waitForSuccessfulExit(): Promise<void> {
@@ -267,7 +284,7 @@ export class TestShell {
   }
 
   async waitForPromptOrExit(
-    opts: { timeout?: number; start?: number } = {}
+    opts: TestShellWaitForPromptOptions & { start?: number } = {}
   ): Promise<TestShellStartupResult> {
     return Promise.race([
       this.waitForPrompt(opts.start ?? 0, opts).then(
@@ -283,21 +300,29 @@ export class TestShell {
     this._process.kill(signal);
   }
 
-  writeInput(chars: string, { end = false } = {}): void {
+  writeInput(
+    chars: string,
+    {
+      end = false,
+      requireFinishedInitialization = true,
+    }: TestShellInputOptions = {}
+  ): void {
+    if (requireFinishedInitialization && !this._initializationKnownToBeFinished)
+      throw new Error('Wait for shell to be initialized before writing input');
     this._process.stdin.write(chars);
     if (end) this._process.stdin.end();
   }
 
-  writeInputLine(chars: string): void {
-    this.writeInput(`${chars}\n`);
+  writeInputLine(chars: string, options?: TestShellInputOptions): void {
+    this.writeInput(`${chars}\n`, options);
   }
 
   async executeLine(
     line: string,
-    opts: { timeout?: number; promptPattern?: RegExp } = {}
+    opts: TestShellWaitForPromptOptions & TestShellInputOptions = {}
   ): Promise<string> {
     const previousOutputLength = this._output.length;
-    this.writeInputLine(line);
+    this.writeInputLine(line, opts);
     await this.waitForPrompt(previousOutputLength, opts);
     return this._output.slice(previousOutputLength);
   }