From 42c456eb7a57ebcd4061940c71e7525984106b02 Mon Sep 17 00:00:00 2001
From: Bob Olde Hampsink <b.oldehampsink@itmundi.nl>
Date: Tue, 20 Oct 2015 14:31:10 +0200
Subject: [PATCH 1/9] Run migrations on plugin update

---
 README.md                              |  1 +
 services/Schematic_PluginsService.php  | 40 ++++++++++++++++-------
 tests/Schematic_PluginsServiceTest.php | 45 +++++++++++++++++++-------
 3 files changed, 63 insertions(+), 23 deletions(-)

diff --git a/README.md b/README.md
index 708db475..3cf6bfec 100644
--- a/README.md
+++ b/README.md
@@ -22,6 +22,7 @@ You can also generate a schema.yml with
 ```
 
 ## Hooks
+
 * Has a hook "registerMigrationService" to add exports with your own data.
 
 ```php
diff --git a/services/Schematic_PluginsService.php b/services/Schematic_PluginsService.php
index e59e0d4e..74b92229 100644
--- a/services/Schematic_PluginsService.php
+++ b/services/Schematic_PluginsService.php
@@ -20,9 +20,8 @@ class Schematic_PluginsService extends BaseApplicationComponent
      */
     protected $resultModel;
 
-
     /**
-     * Constructor to setup result model
+     * Constructor to setup result model.
      */
     public function __construct()
     {
@@ -38,7 +37,16 @@ protected function getPluginService()
     }
 
     /**
-     * Installs plugin by handle
+     * @return MigrationsService
+     */
+    protected function getMigrationsService()
+    {
+        return craft()->migrations;
+    }
+
+    /**
+     * Installs plugin by handle.
+     *
      * @param string $handle
      */
     protected function installPluginByHandle($handle)
@@ -51,7 +59,8 @@ protected function installPluginByHandle($handle)
     }
 
     /**
-     * Uninstalls plugin by handle
+     * Uninstalls plugin by handle.
+     *
      * @param $handle
      */
     protected function uninstallPluginByHandle($handle)
@@ -60,8 +69,10 @@ protected function uninstallPluginByHandle($handle)
     }
 
     /**
-     * Returns plugin by handle
+     * Returns plugin by handle.
+     *
      * @param string $handle
+     *
      * @return BasePlugin|null
      */
     protected function getPlugin($handle)
@@ -75,9 +86,10 @@ protected function getPlugin($handle)
     }
 
     /**
-     * Toggles plugin based on enabled flag
+     * Toggles plugin based on enabled flag.
+     *
      * @param string $handle
-     * @param bool $isEnabled
+     * @param bool   $isEnabled
      */
     protected function togglePluginByHandle($handle, $isEnabled)
     {
@@ -89,7 +101,8 @@ protected function togglePluginByHandle($handle, $isEnabled)
     }
 
     /**
-     * Add error to errors collection
+     * Add error to errors collection.
+     *
      * @param $message
      */
     private function addError($message)
@@ -99,6 +112,7 @@ private function addError($message)
 
     /**
      * @param BasePlugin $plugin
+     *
      * @return array
      */
     private function getPluginDefinition(BasePlugin $plugin)
@@ -106,26 +120,29 @@ private function getPluginDefinition(BasePlugin $plugin)
         return array(
             'isInstalled'       => $plugin->isInstalled,
             'isEnabled'         => $plugin->isEnabled,
-            'settings'          => $plugin->getSettings()->attributes
+            'settings'          => $plugin->getSettings()->attributes,
         );
     }
 
     /**
      * @param array $pluginDefinitions
+     *
      * @return Schematic_ResultModel
      */
     public function import(array $pluginDefinitions)
     {
         foreach ($pluginDefinitions as $handle => $pluginDefinition) {
-            if($plugin = $this->getPlugin($handle)) {
+            if ($plugin = $this->getPlugin($handle)) {
                 if ($pluginDefinition['isInstalled']) {
                     $this->installPluginByHandle($handle);
 
                     $this->togglePluginByHandle($handle, $pluginDefinition['isEnabled']);
 
-                    if(array_key_exists('settings', $pluginDefinition)){
+                    if (array_key_exists('settings', $pluginDefinition)) {
                         $this->getPluginService()->savePluginSettings($plugin, $pluginDefinition['settings']);
                     }
+
+                    $this->getMigrationsService()->runToTop($plugin);
                 } else {
                     $this->uninstallPluginByHandle($handle);
                 }
@@ -147,6 +164,7 @@ public function export()
             $pluginDefinitions[$handle] = $this->getPluginDefinition($plugin);
         }
         ksort($pluginDefinitions);
+
         return $pluginDefinitions;
     }
 }
diff --git a/tests/Schematic_PluginsServiceTest.php b/tests/Schematic_PluginsServiceTest.php
index 1b0382ad..8bc138e4 100644
--- a/tests/Schematic_PluginsServiceTest.php
+++ b/tests/Schematic_PluginsServiceTest.php
@@ -5,7 +5,7 @@
 use PHPUnit_Framework_MockObject_MockObject as Mock;
 
 /**
- * Class Schematic_PluginsServiceTest
+ * Class Schematic_PluginsServiceTest.
  *
  * @author    Itmundi
  * @copyright Copyright (c) 2015, Itmundi
@@ -49,6 +49,7 @@ public static function setUpBeforeClass()
      * @param bool $enablePluginResponse
      * @param bool $disablePluginResponse
      * @param bool $uninstallPluginResponse
+     *
      * @return PluginsService|Mock
      */
     public function getMockPluginsService(
@@ -62,7 +63,7 @@ public function getMockPluginsService(
 
         $mock->expects($this->any())->method('getPlugin')->willReturn(($returnPlugin) ? $this->getMockBasePlugin() : null);
 
-        if($installPluginResponse) {
+        if ($installPluginResponse) {
             $mock->expects($this->any())->method('installPlugin')->willReturn($installPluginResponse);
         } else {
             $mock->expects($this->any())->method('installPlugin')->willThrowException(new Exception());
@@ -75,6 +76,17 @@ public function getMockPluginsService(
         return $mock;
     }
 
+    /**
+     * @return MigrationsService|Mock
+     */
+    public function getMockMigrationsService()
+    {
+        $mock = $this->getMockBuilder('Craft\MigrationsService')->getMock();
+        $mock->expects($this->any())->method('runToTop')->willReturn(true);
+
+        return $mock;
+    }
+
     /**
      * @return Mock|BasePlugin
      */
@@ -86,7 +98,8 @@ public function getMockBasePlugin()
     }
 
     /**
-     * Test default import functionality
+     * Test default import functionality.
+     *
      * @covers ::import
      */
     public function testImportWithInstalledPlugins()
@@ -94,6 +107,8 @@ public function testImportWithInstalledPlugins()
         $data = $this->getPluginsData();
         $mockPluginsService = $this->getMockPluginsService();
         $this->setComponent(craft(), 'plugins', $mockPluginsService);
+        $mockMigrationsService = $this->getMockMigrationsService();
+        $this->setComponent(craft(), 'migrations', $mockMigrationsService);
 
         $import = $this->schematicPluginsService->import($data);
 
@@ -102,7 +117,8 @@ public function testImportWithInstalledPlugins()
     }
 
     /**
-     * Test default import functionality
+     * Test default import functionality.
+     *
      * @covers ::import
      */
     public function testImportWithInstalledDisabledPlugins()
@@ -119,7 +135,8 @@ public function testImportWithInstalledDisabledPlugins()
     }
 
     /**
-     * Test default import functionality
+     * Test default import functionality.
+     *
      * @covers ::import
      */
     public function testImportWithMissingPlugin()
@@ -136,7 +153,8 @@ public function testImportWithMissingPlugin()
     }
 
     /**
-     * Test default import functionality
+     * Test default import functionality.
+     *
      * @covers ::import
      */
     public function testImportWithInstallException()
@@ -153,7 +171,8 @@ public function testImportWithInstallException()
     }
 
     /**
-     * Test default import functionality
+     * Test default import functionality.
+     *
      * @covers ::import
      */
     public function testImportWithNotInstalledPlugin()
@@ -171,7 +190,8 @@ public function testImportWithNotInstalledPlugin()
     }
 
     /**
-     * Test export functionality
+     * Test export functionality.
+     *
      * @covers ::export
      */
     public function testExport()
@@ -198,7 +218,8 @@ public function testExport()
     }
 
     /**
-     * Returns plugins data
+     * Returns plugins data.
+     *
      * @return array
      */
     public function getPluginsData()
@@ -210,9 +231,9 @@ public function getPluginsData()
                 'settings'          => array(
                     'pluginName'    => 'Menu',
                     'canDoActions'  => '',
-                    'quietErrors'   => ''
-                )
-            )
+                    'quietErrors'   => '',
+                ),
+            ),
         );
     }
 }

From c9fc96d895268fdf94ee548350b6fd0493d4ab5b Mon Sep 17 00:00:00 2001
From: Bob Olde Hampsink <b.oldehampsink@itmundi.nl>
Date: Tue, 20 Oct 2015 14:40:34 +0200
Subject: [PATCH 2/9] Fix mocks

---
 tests/Schematic_PluginsServiceTest.php | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tests/Schematic_PluginsServiceTest.php b/tests/Schematic_PluginsServiceTest.php
index 8bc138e4..262db4df 100644
--- a/tests/Schematic_PluginsServiceTest.php
+++ b/tests/Schematic_PluginsServiceTest.php
@@ -127,6 +127,8 @@ public function testImportWithInstalledDisabledPlugins()
         $data['itmundiplugin']['isEnabled'] = false;
         $mockPluginsService = $this->getMockPluginsService();
         $this->setComponent(craft(), 'plugins', $mockPluginsService);
+        $mockMigrationsService = $this->getMockMigrationsService();
+        $this->setComponent(craft(), 'migrations', $mockMigrationsService);
 
         $import = $this->schematicPluginsService->import($data);
 
@@ -163,6 +165,8 @@ public function testImportWithInstallException()
 
         $mockPluginsService = $this->getMockPluginsService(true, false);
         $this->setComponent(craft(), 'plugins', $mockPluginsService);
+        $mockMigrationsService = $this->getMockMigrationsService();
+        $this->setComponent(craft(), 'migrations', $mockMigrationsService);
 
         $import = $this->schematicPluginsService->import($data);
 

From 36da3b38d431a834c003810a8027bc282565553e Mon Sep 17 00:00:00 2001
From: Bob Olde Hampsink <b.oldehampsink@itmundi.nl>
Date: Thu, 22 Oct 2015 10:55:33 +0200
Subject: [PATCH 3/9] Log to both cli and logs, more verbose

---
 consolecommands/SchematicCommand.php | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/consolecommands/SchematicCommand.php b/consolecommands/SchematicCommand.php
index e9876aa1..0b9a271a 100644
--- a/consolecommands/SchematicCommand.php
+++ b/consolecommands/SchematicCommand.php
@@ -25,18 +25,17 @@ public function actionImport($file = 'craft/config/schema.yml', $force = false)
     {
         if (!IOHelper::fileExists($file)) {
             $this->usageError(Craft::t('File not found.'));
-            exit(1);
         }
 
         $result = craft()->schematic->importFromYaml($file, $force);
 
         if (!$result->hasErrors()) {
-            echo Craft::t('Loaded schema from {file}', array('file' => $file))."\n";
+            SchematicPlugin::log(Craft::t('Loaded schema from {file}', array('file' => $file)));
             exit(0);
         }
 
-        echo Craft::t('There was an error loading schema from {file}', array('file' => $file))."\n";
-        print_r($result->errors);
+        SchematicPlugin::log(Craft::t('There was an error loading schema from {file}', array('file' => $file)));
+        print_r($result->getErrors());
         exit(1);
     }
 
@@ -49,7 +48,7 @@ public function actionExport($file = 'craft/config/schema.yml')
     {
         craft()->schematic->exportToYaml($file);
 
-        echo Craft::t('Exported schema to {file}', array('file' => $file))."\n";
+        SchematicPlugin::log(Craft::t('Exported schema to {file}', array('file' => $file)));
         exit(0);
     }
 }

From 78ca6e9097d419dd0187a5a6aae2a7504ceec5fc Mon Sep 17 00:00:00 2001
From: Bob Olde Hampsink <b.oldehampsink@itmundi.nl>
Date: Thu, 22 Oct 2015 11:32:20 +0200
Subject: [PATCH 4/9] Run migrations before saving settings

---
 services/Schematic_PluginsService.php | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/services/Schematic_PluginsService.php b/services/Schematic_PluginsService.php
index 8c8a05bf..97f50f37 100644
--- a/services/Schematic_PluginsService.php
+++ b/services/Schematic_PluginsService.php
@@ -116,11 +116,11 @@ public function import(array $pluginDefinitions, $force = false)
 
                     $this->togglePluginByHandle($handle, $pluginDefinition['isEnabled']);
 
+                    $this->getMigrationsService()->runToTop($plugin);
+
                     if (array_key_exists('settings', $pluginDefinition)) {
                         $this->getPluginService()->savePluginSettings($plugin, $pluginDefinition['settings']);
                     }
-
-                    $this->getMigrationsService()->runToTop($plugin);
                 } else {
                     $this->uninstallPluginByHandle($handle);
                 }

From d29a8016b7ac9e0fd4cc627d1f9c6f6afa5fa227 Mon Sep 17 00:00:00 2001
From: Bob Olde Hampsink <b.oldehampsink@itmundi.nl>
Date: Thu, 22 Oct 2015 11:33:00 +0200
Subject: [PATCH 5/9] Return exit code instead of exiting

---
 consolecommands/SchematicCommand.php | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/consolecommands/SchematicCommand.php b/consolecommands/SchematicCommand.php
index 0b9a271a..1bc94afc 100644
--- a/consolecommands/SchematicCommand.php
+++ b/consolecommands/SchematicCommand.php
@@ -31,12 +31,14 @@ public function actionImport($file = 'craft/config/schema.yml', $force = false)
 
         if (!$result->hasErrors()) {
             SchematicPlugin::log(Craft::t('Loaded schema from {file}', array('file' => $file)));
-            exit(0);
+
+            return 0;
         }
 
         SchematicPlugin::log(Craft::t('There was an error loading schema from {file}', array('file' => $file)));
         print_r($result->getErrors());
-        exit(1);
+
+        return 1;
     }
 
     /**
@@ -49,6 +51,7 @@ public function actionExport($file = 'craft/config/schema.yml')
         craft()->schematic->exportToYaml($file);
 
         SchematicPlugin::log(Craft::t('Exported schema to {file}', array('file' => $file)));
-        exit(0);
+
+        return 0;
     }
 }

From 44abee74118b039fe537cafbcca1905eec5411e8 Mon Sep 17 00:00:00 2001
From: Bob Olde Hampsink <b.oldehampsink@itmundi.nl>
Date: Thu, 22 Oct 2015 11:40:55 +0200
Subject: [PATCH 6/9] Apply new plugin info after migrations

---
 services/Schematic_PluginsService.php  | 25 ++++++++++++++++++++++++-
 tests/Schematic_PluginsServiceTest.php | 17 +++++++++++++++++
 2 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/services/Schematic_PluginsService.php b/services/Schematic_PluginsService.php
index 97f50f37..0a31eb1a 100644
--- a/services/Schematic_PluginsService.php
+++ b/services/Schematic_PluginsService.php
@@ -31,6 +31,14 @@ protected function getMigrationsService()
         return craft()->migrations;
     }
 
+    /**
+     * @return UpdatesService
+     */
+    protected function getUpdatesService()
+    {
+        return craft()->updates;
+    }
+
     /**
      * Installs plugin by handle.
      *
@@ -87,6 +95,21 @@ protected function togglePluginByHandle($handle, $isEnabled)
         }
     }
 
+    /**
+     * Run plugin migrations automatically.
+     *
+     * @param BasePlugin $plugin
+     */
+    protected function runMigrations(BasePlugin $plugin)
+    {
+        if (!$this->getMigrationsService()->runToTop($plugin)) {
+            throw new Exception(Craft::t('There was a problem updating your database.'));
+        }
+        if (!$this->getUpdatesService()->setNewPluginInfo($plugin)) {
+            throw new Exception(Craft::t('The update was performed successfully, but there was a problem setting the new info in the plugins table.'));
+        }
+    }
+
     /**
      * @param BasePlugin $plugin
      *
@@ -116,7 +139,7 @@ public function import(array $pluginDefinitions, $force = false)
 
                     $this->togglePluginByHandle($handle, $pluginDefinition['isEnabled']);
 
-                    $this->getMigrationsService()->runToTop($plugin);
+                    $this->runMigrations($plugin);
 
                     if (array_key_exists('settings', $pluginDefinition)) {
                         $this->getPluginService()->savePluginSettings($plugin, $pluginDefinition['settings']);
diff --git a/tests/Schematic_PluginsServiceTest.php b/tests/Schematic_PluginsServiceTest.php
index 4c4fc960..6f1febd7 100644
--- a/tests/Schematic_PluginsServiceTest.php
+++ b/tests/Schematic_PluginsServiceTest.php
@@ -88,6 +88,17 @@ public function getMockMigrationsService()
         return $mock;
     }
 
+    /**
+     * @return UpdatesService|Mock
+     */
+    public function getMockUpdatesService()
+    {
+        $mock = $this->getMockBuilder('Craft\UpdatesService')->getMock();
+        $mock->expects($this->any())->method('setNewPluginInfo')->willReturn(true);
+
+        return $mock;
+    }
+
     /**
      * @return Mock|BasePlugin
      */
@@ -110,6 +121,8 @@ public function testImportWithInstalledPlugins()
         $this->setComponent(craft(), 'plugins', $mockPluginsService);
         $mockMigrationsService = $this->getMockMigrationsService();
         $this->setComponent(craft(), 'migrations', $mockMigrationsService);
+        $mockUpdatesService = $this->getMockUpdatesService();
+        $this->setComponent(craft(), 'updates', $mockUpdatesService);
 
         $import = $this->schematicPluginsService->import($data);
 
@@ -130,6 +143,8 @@ public function testImportWithInstalledDisabledPlugins()
         $this->setComponent(craft(), 'plugins', $mockPluginsService);
         $mockMigrationsService = $this->getMockMigrationsService();
         $this->setComponent(craft(), 'migrations', $mockMigrationsService);
+        $mockUpdatesService = $this->getMockUpdatesService();
+        $this->setComponent(craft(), 'updates', $mockUpdatesService);
 
         $import = $this->schematicPluginsService->import($data);
 
@@ -168,6 +183,8 @@ public function testImportWithInstallException()
         $this->setComponent(craft(), 'plugins', $mockPluginsService);
         $mockMigrationsService = $this->getMockMigrationsService();
         $this->setComponent(craft(), 'migrations', $mockMigrationsService);
+        $mockUpdatesService = $this->getMockUpdatesService();
+        $this->setComponent(craft(), 'updates', $mockUpdatesService);
 
         $import = $this->schematicPluginsService->import($data);
 

From 80449eea791459ea0edfba91b5504bb4ebb7ee37 Mon Sep 17 00:00:00 2001
From: Bob Olde Hampsink <b.oldehampsink@itmundi.nl>
Date: Thu, 22 Oct 2015 11:59:21 +0200
Subject: [PATCH 7/9] Prevent code duplication in unit test

---
 tests/Schematic_PluginsServiceTest.php | 36 +++++++++++++-------------
 1 file changed, 18 insertions(+), 18 deletions(-)

diff --git a/tests/Schematic_PluginsServiceTest.php b/tests/Schematic_PluginsServiceTest.php
index 6f1febd7..45559321 100644
--- a/tests/Schematic_PluginsServiceTest.php
+++ b/tests/Schematic_PluginsServiceTest.php
@@ -44,6 +44,19 @@ public static function setUpBeforeClass()
         require_once __DIR__.'/../services/Schematic_PluginsService.php';
     }
 
+    /**
+     * Prevent code duplication by mocking multiple services.
+     */
+    public function mockMultipleServices()
+    {
+        $mockPluginsService = $this->getMockPluginsService();
+        $this->setComponent(craft(), 'plugins', $mockPluginsService);
+        $mockMigrationsService = $this->getMockMigrationsService();
+        $this->setComponent(craft(), 'migrations', $mockMigrationsService);
+        $mockUpdatesService = $this->getMockUpdatesService();
+        $this->setComponent(craft(), 'updates', $mockUpdatesService);
+    }
+
     /**
      * @param bool $returnPlugin
      * @param bool $installPluginResponse
@@ -117,12 +130,8 @@ public function getMockBasePlugin()
     public function testImportWithInstalledPlugins()
     {
         $data = $this->getPluginsData();
-        $mockPluginsService = $this->getMockPluginsService();
-        $this->setComponent(craft(), 'plugins', $mockPluginsService);
-        $mockMigrationsService = $this->getMockMigrationsService();
-        $this->setComponent(craft(), 'migrations', $mockMigrationsService);
-        $mockUpdatesService = $this->getMockUpdatesService();
-        $this->setComponent(craft(), 'updates', $mockUpdatesService);
+
+        $this->mockMultipleServices();
 
         $import = $this->schematicPluginsService->import($data);
 
@@ -139,12 +148,8 @@ public function testImportWithInstalledDisabledPlugins()
     {
         $data = $this->getPluginsData();
         $data['itmundiplugin']['isEnabled'] = false;
-        $mockPluginsService = $this->getMockPluginsService();
-        $this->setComponent(craft(), 'plugins', $mockPluginsService);
-        $mockMigrationsService = $this->getMockMigrationsService();
-        $this->setComponent(craft(), 'migrations', $mockMigrationsService);
-        $mockUpdatesService = $this->getMockUpdatesService();
-        $this->setComponent(craft(), 'updates', $mockUpdatesService);
+
+        $this->mockMultipleServices();
 
         $import = $this->schematicPluginsService->import($data);
 
@@ -179,12 +184,7 @@ public function testImportWithInstallException()
     {
         $data = $this->getPluginsData();
 
-        $mockPluginsService = $this->getMockPluginsService(true, false);
-        $this->setComponent(craft(), 'plugins', $mockPluginsService);
-        $mockMigrationsService = $this->getMockMigrationsService();
-        $this->setComponent(craft(), 'migrations', $mockMigrationsService);
-        $mockUpdatesService = $this->getMockUpdatesService();
-        $this->setComponent(craft(), 'updates', $mockUpdatesService);
+        $this->mockMultipleServices();
 
         $import = $this->schematicPluginsService->import($data);
 

From 387485fc39abb1c89d4cf9b3702c9f874349b459 Mon Sep 17 00:00:00 2001
From: Bob Olde Hampsink <b.oldehampsink@itmundi.nl>
Date: Thu, 22 Oct 2015 12:04:12 +0200
Subject: [PATCH 8/9] Install when not installed, migrate when installed
 already

---
 services/Schematic_PluginsService.php | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/services/Schematic_PluginsService.php b/services/Schematic_PluginsService.php
index 0a31eb1a..98f794b2 100644
--- a/services/Schematic_PluginsService.php
+++ b/services/Schematic_PluginsService.php
@@ -135,12 +135,14 @@ public function import(array $pluginDefinitions, $force = false)
         foreach ($pluginDefinitions as $handle => $pluginDefinition) {
             if ($plugin = $this->getPlugin($handle)) {
                 if ($pluginDefinition['isInstalled']) {
-                    $this->installPluginByHandle($handle);
+                    if (!$plugin->isInstalled) {
+                        $this->installPluginByHandle($handle);
+                    } else {
+                        $this->runMigrations($plugin);
+                    }
 
                     $this->togglePluginByHandle($handle, $pluginDefinition['isEnabled']);
 
-                    $this->runMigrations($plugin);
-
                     if (array_key_exists('settings', $pluginDefinition)) {
                         $this->getPluginService()->savePluginSettings($plugin, $pluginDefinition['settings']);
                     }

From 7904cdf10df86ee22a281ed624b6c25f658418db Mon Sep 17 00:00:00 2001
From: Bob Olde Hampsink <b.oldehampsink@itmundi.nl>
Date: Thu, 22 Oct 2015 12:44:51 +0200
Subject: [PATCH 9/9] Fix failing unit tests

---
 tests/Schematic_PluginsServiceTest.php | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/tests/Schematic_PluginsServiceTest.php b/tests/Schematic_PluginsServiceTest.php
index 45559321..16aebbe5 100644
--- a/tests/Schematic_PluginsServiceTest.php
+++ b/tests/Schematic_PluginsServiceTest.php
@@ -46,10 +46,15 @@ public static function setUpBeforeClass()
 
     /**
      * Prevent code duplication by mocking multiple services.
+     *
+     * @param bool $returnPlugin
+     * @param bool $installPluginResponse
      */
-    public function mockMultipleServices()
-    {
-        $mockPluginsService = $this->getMockPluginsService();
+    public function mockMultipleServices(
+        $returnPlugin = true,
+        $installPluginResponse = true
+    ) {
+        $mockPluginsService = $this->getMockPluginsService($returnPlugin, $installPluginResponse);
         $this->setComponent(craft(), 'plugins', $mockPluginsService);
         $mockMigrationsService = $this->getMockMigrationsService();
         $this->setComponent(craft(), 'migrations', $mockMigrationsService);
@@ -184,7 +189,7 @@ public function testImportWithInstallException()
     {
         $data = $this->getPluginsData();
 
-        $this->mockMultipleServices();
+        $this->mockMultipleServices(true, false);
 
         $import = $this->schematicPluginsService->import($data);