diff --git a/src/MigrateBatchExecutable.php b/src/MigrateBatchExecutable.php
index 85c28b6496478997e9de192ec193c884203e913b..13db168c46edefe8fe9e24e11a717afca37cfefc 100644
--- a/src/MigrateBatchExecutable.php
+++ b/src/MigrateBatchExecutable.php
@@ -80,10 +80,10 @@ class MigrateBatchExecutable extends BaseMigrateBatchExecutable {
    *   The migration id.
    * @param array $options
    *   The batch executable options.
-   * @param array $context
+   * @param array|\DrushBatchContext $context
    *   The sandbox context.
    */
-  public static function batchProcessImport($migration_id, array $options, array &$context) {
+  public static function batchProcessImport($migration_id, array $options, &$context) {
     if (empty($context['sandbox'])) {
       $context['finished'] = 0;
       $context['sandbox'] = [];