From c6e34935e85e176754108400ed2a6bf0bb5f6a83 Mon Sep 17 00:00:00 2001
From: Marcelo Shima <marceloshima@gmail.com>
Date: Wed, 24 Feb 2021 15:50:15 -0300
Subject: [PATCH] Make a mockedGenerator a Generator child.

---
 lib/index.js | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/lib/index.js b/lib/index.js
index d3402b5..1406b53 100644
--- a/lib/index.js
+++ b/lib/index.js
@@ -8,6 +8,7 @@
 const fs = require('fs');
 const path = require('path');
 const sinon = require('sinon');
+const Generator = require('yeoman-generator');
 const adapter = require('./adapter');
 
 /**
@@ -143,11 +144,13 @@ exports.mockLocalConfig = function (generator, localConfig) {
  * Create a mocked generator
  */
 
-exports.createMockedGenerator = () => {
-  const generator = sinon.stub();
-  generator.prototype.run = sinon.stub();
-  generator.prototype.runWithOptions = sinon.stub();
-  generator.prototype.queueOwnTasks = sinon.stub();
+exports.createMockedGenerator = (constructor = Generator) => {
+  const generator = sinon.spy(constructor);
+  ['run', 'queueTasks', 'runWithOptions', 'queueOwnTasks'].forEach(methodName => {
+    if (Generator.prototype[methodName]) {
+      generator.prototype[methodName] = sinon.stub();
+    }
+  });
   return generator;
 };