add withEnvironmentRun helper (#241)

Author: mshima

src/run-context.ts

@@ -103,6 +103,7 @@ export class RunContextBase<GeneratorType extends BaseGenerator = DefaultGenerat
   private ran = false;
   private errored = false;
   private readonly beforePrepareCallbacks: Array<(this: this) => void | Promise<void>> = [];
+  private environmentRun?: (this: this, env: DefaultEnvironmentApi, gen: GeneratorType) => void;
 
   /**
    * This class provide a run context object to façade the complexity involved in setting
@@ -151,7 +152,8 @@ export class RunContextBase<GeneratorType extends BaseGenerator = DefaultGenerat
     }
 
     try {
-      await this.env.runGenerator(this.generator as any);
+      const environmentRun = this.environmentRun ?? ((env, generator) => env.runGenerator(generator));
+      await environmentRun.call(this, this.env, this.generator);
     } finally {
       this.helpers.restorePrompt(this.env);
       this.completed = true;
@@ -301,6 +303,17 @@ export class RunContextBase<GeneratorType extends BaseGenerator = DefaultGenerat
     return this;
   }
 
+  /**
+   * Customize enviroment run method.
+   *
+   * @param callback
+   * @return {this} run context instance
+   */
+  withEnvironmentRun(callback: (this: this, env: DefaultEnvironmentApi, gen: GeneratorType) => void) {
+    this.environmentRun = callback;
+    return this;
+  }
+
   /**
    * Run lookup on the environment.
    *

test/run-context.spec.ts

@@ -6,12 +6,13 @@ import process from 'node:process';
 import { createRequire } from 'node:module';
 import { mock } from 'node:test';
 import { promisify as promisify_ } from 'node:util';
-import { afterEach, beforeEach, describe, expect, it } from 'vitest';
+import { afterEach, beforeAll, beforeEach, describe, expect, it } from 'vitest';
 import Generator from 'yeoman-generator';
 import tempDirectory from 'temp-dir';
 import { RunContextBase as RunContext } from '../src/run-context.js';
 import helpers from '../src/helpers.js';
 import { DummyPrompt } from '../src/adapter.js';
+import { BaseEnvironmentOptions } from '@yeoman/types';
 
 /* Remove argument from promisify return */
 const promisify = function_ => () => promisify_(function_)();
@@ -21,7 +22,7 @@ const __dirname = dirname(fileURLToPath(import.meta.url));
 const tmpdir = path.join(tempDirectory, 'yeoman-run-context');
 
 describe('RunContext', () => {
-  const environmentOptions = { foo: 'bar' };
+  let environmentOptions: BaseEnvironmentOptions | undefined;
   let context: RunContext;
   let execSpy;
   let Dummy;
@@ -61,6 +62,9 @@ describe('RunContext', () => {
   );
 
   describe('constructor', () => {
+    beforeAll(() => {
+      environmentOptions = { foo: 'bar' };
+    });
     it(
       'forwards envOptions to the environment',
       promisify(done => {
@@ -856,6 +860,32 @@ describe('RunContext', () => {
     );
   });
 
+  describe('#withEnvironmentRun()', () => {
+    it('calls runGenerator by default', async () => {
+      let mockedRunGenerator: ReturnType<typeof mock.fn>;
+      await context
+        .withEnvironment(environment => {
+          mockedRunGenerator = mock.method(environment, 'runGenerator');
+        })
+        .toPromise();
+      expect(mockedRunGenerator!.mock.callCount()).toBe(1);
+    });
+
+    it('calls custom environment run method', async () => {
+      let mockedRunGenerator: ReturnType<typeof mock.fn>;
+      const mockedEnvironmentRun = mock.fn();
+      await context
+        .withEnvironment(environment => {
+          mockedRunGenerator = mock.method(environment, 'runGenerator');
+        })
+        .withEnvironmentRun(mockedEnvironmentRun)
+        .toPromise();
+
+      expect(mockedRunGenerator!.mock.callCount()).toBe(0);
+      expect(mockedEnvironmentRun!.mock.callCount()).toBe(1);
+    });
+  });
+
   describe('#withLocalConfig()', () => {
     it(
       'provides config to the generator',