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;
 };