diff --git a/drush.services.yml b/drush.services.yml
index e1a43c7..ba6b52d 100644
--- a/drush.services.yml
+++ b/drush.services.yml
@@ -1,6 +1,6 @@
 services:
   migrate_tools.commands:
     class: \Drupal\migrate_tools\Commands\MigrateToolsCommands
-    arguments: ['@plugin.manager.migration', '@date.formatter']
+    arguments: ['@plugin.manager.migration', '@date.formatter', '@entity_type.manager']
     tags:
       - { name: drush.command }
diff --git a/src/Commands/MigrateToolsCommands.php b/src/Commands/MigrateToolsCommands.php
index 8233ce3..65fb4d2 100644
--- a/src/Commands/MigrateToolsCommands.php
+++ b/src/Commands/MigrateToolsCommands.php
@@ -5,11 +5,11 @@ namespace Drupal\migrate_tools\Commands;
 use Consolidation\OutputFormatters\StructuredData\RowsOfFields;
 use Drupal\Component\Utility\Unicode;
 use Drupal\Core\Datetime\DateFormatter;
+use Drupal\Core\Entity\EntityTypeManagerInterface;
 use Drupal\migrate\Exception\RequirementsException;
 use Drupal\migrate\Plugin\MigrationInterface;
 use Drupal\migrate\Plugin\MigrationPluginManager;
 use Drupal\migrate\Plugin\RequirementsInterface;
-use Drupal\migrate_plus\Entity\MigrationGroup;
 use Drupal\migrate_tools\Drush9LogMigrateMessage;
 use Drupal\migrate_tools\MigrateExecutable;
 use Drush\Commands\DrushCommands;
@@ -34,19 +34,30 @@ class MigrateToolsCommands extends DrushCommands {
   protected $dateFormatter;
 
   /**
+   * Entity type manager.
+   *
+   * @var \Drupal\Core\Entity\EntityTypeManagerInterface
+   */
+  protected $entityTypeManager;
+
+  /**
    * MigrateToolsCommands constructor.
    *
    * @param \Drupal\migrate\Plugin\MigrationPluginManager $migrationPluginManager
    *   Migration Plugin Manager service.
    * @param \Drupal\Core\Datetime\DateFormatter $dateFormatter
    *   Date formatter service.
+   * @param \Drupal\Core\Entity\EntityTypeManagerInterface $entityTypeManager
+   *   Entity type manager service.
    */
   public function __construct(
     MigrationPluginManager $migrationPluginManager,
-    DateFormatter $dateFormatter
+    DateFormatter $dateFormatter,
+    EntityTypeManagerInterface $entityTypeManager
   ) {
     $this->migrationPluginManager = $migrationPluginManager;
     $this->dateFormatter = $dateFormatter;
+    $this->entityTypeManager = $entityTypeManager;
   }
 
   /**
@@ -106,7 +117,8 @@ class MigrateToolsCommands extends DrushCommands {
     $table = [];
     // Take it one group at a time, listing the migrations within each group.
     foreach ($migrations as $group_id => $migration_list) {
-      $group = MigrationGroup::load($group_id);
+      /** @var \Drupal\migrate_plus\Entity\MigrationGroup $group */
+      $group = $this->entityTypeManager->getStorage('migration_group')->load($group_id);
       $group_name = !empty($group) ? "{$group->label()} ({$group->id()})" : $group_id;
 
       $this->output()->writeln(dt('Group: @name', ['@name' => $group_name]));
@@ -274,11 +286,7 @@ class MigrateToolsCommands extends DrushCommands {
     // Sort the matched migrations by group.
     if (!empty($matched_migrations)) {
       foreach ($matched_migrations as $id => $migration) {
-        $configured_group_id = empty(
-        $migration->get(
-          'migration_group'
-        )
-        ) ? 'default' : $migration->get('migration_group');
+        $configured_group_id = empty($migration->get('migration_group')) ? 'default' : $migration->get('migration_group');
         $migrations[$configured_group_id][$id] = $migration;
       }
     }
@@ -348,11 +356,7 @@ class MigrateToolsCommands extends DrushCommands {
     $all = $options['all'];
     $additional_options = [];
     if (!$all && !$group_names && !$migration_names && !$tag_names) {
-      throw new \Exception(
-        dt(
-          'You must specify --all, --group, --tag or one or more migration names separated by commas'
-        )
-      );
+      throw new \Exception(dt('You must specify --all, --group, --tag or one or more migration names separated by commas'));
     }
 
     foreach (['limit', 'feedback', 'idlist', 'update', 'force'] as $option) {
@@ -404,6 +408,9 @@ class MigrateToolsCommands extends DrushCommands {
    * @validate-module-enabled migrate_tools
    *
    * @aliases mr, migrate-rollback
+   *
+   * @throws \Exception
+   *   If there are not enough parameters to the command.
    */
   public function rollback(
     $migration_names = '',
@@ -419,13 +426,7 @@ class MigrateToolsCommands extends DrushCommands {
     $all = $options['all'];
     $additional_options = [];
     if (!$all && !$group_names && !$migration_names && !$tag_names) {
-      drush_set_error(
-        'MIGRATE_ERROR',
-        dt(
-          'You must specify --all, --group, --tag, or one or more migration names separated by commas'
-        )
-      );
-      return;
+      throw new \Exception(dt('You must specify --all, --group, --tag, or one or more migration names separated by commas'));
     }
 
     if ($options['feedback']) {