diff --git a/core/modules/migrate/tests/src/Unit/MigrateSqlSourceTestCase.php b/core/modules/migrate/tests/src/Unit/MigrateSqlSourceTestCase.php
index 16d7e0b..70c544a 100644
--- a/core/modules/migrate/tests/src/Unit/MigrateSqlSourceTestCase.php
+++ b/core/modules/migrate/tests/src/Unit/MigrateSqlSourceTestCase.php
@@ -76,11 +76,24 @@ protected function setUp() {
     $migration->expects($this->any())
       ->method('getHighWater')
       ->will($this->returnValue(static::ORIGINAL_HIGH_WATER));
-    // Need the test class, not the original because we need a setDatabase method. This is not pretty :/
-    $plugin_class  = preg_replace('/^Drupal\\\\(\w+)\\\\Plugin\\\\migrate(\\\\source(\\\\.+)?\\\\)([^\\\\]+)$/', 'Drupal\\Tests\\\$1\\Unit$2Test$4', static::PLUGIN_CLASS);
+
+    // Give each test the oppurtunity to provide a test version of the plugin
+    // we're testing.
+    $test_plugin_class = $plugin_class  = preg_replace('/^Drupal\\\\(\w+)\\\\Plugin\\\\migrate(\\\\source(\\\\.+)?\\\\)([^\\\\]+)$/', 'Drupal\\Tests\\\$1\\Unit$2Test$4', static::PLUGIN_CLASS);
+    $plugin_class = class_exists($test_plugin_class) ? $test_plugin_class : static::PLUGIN_CLASS;
     $plugin = new $plugin_class($this->migrationConfiguration['source'], $this->migrationConfiguration['source']['plugin'], array(), $migration, $entity_manager);
-    $plugin->setDatabase($this->getDatabase($this->databaseContents + array('test_map' => array())));
-    $plugin->setModuleHandler($module_handler);
+
+    // Do some reflection to set the database and moduleHandler.
+    $plugin_reflection = new \ReflectionClass($plugin);
+    $database_property = $plugin_reflection->getProperty('database');
+    $database_property->setAccessible(TRUE);
+    $module_handler_property = $plugin_reflection->getProperty('moduleHandler');
+    $module_handler_property->setAccessible(TRUE);
+
+    // Set the database and the module handler onto our plugin.
+    $database_property->setValue($plugin, $this->getDatabase($this->databaseContents + array('test_map' => array())));
+    $module_handler_property->setValue($plugin, $module_handler);
+
     $plugin->setStringTranslation($this->getStringTranslationStub());
     $migration->expects($this->any())
       ->method('getSourcePlugin')
diff --git a/core/modules/migrate_drupal/tests/src/Unit/source/d7/FileTest.php b/core/modules/migrate_drupal/tests/src/Unit/source/d7/FileTest.php
index b473921..66510ab 100644
--- a/core/modules/migrate_drupal/tests/src/Unit/source/d7/FileTest.php
+++ b/core/modules/migrate_drupal/tests/src/Unit/source/d7/FileTest.php
@@ -51,18 +51,3 @@ protected function setUp() {
   }
 
 }
-
-namespace Drupal\Tests\migrate_drupal\Unit\source\d7;
-
-use Drupal\Core\Database\Connection;
-use Drupal\Core\Extension\ModuleHandlerInterface;
-use Drupal\migrate_drupal\Plugin\migrate\source\d7\File;
-
-class TestFile extends File {
-  public function setDatabase(Connection $database) {
-    $this->database = $database;
-  }
-  public function setModuleHandler(ModuleHandlerInterface $module_handler) {
-    $this->moduleHandler = $module_handler;
-  }
-}