Feature allow custom batch processors

includes/functions.php

@@ -16,40 +16,71 @@
  * Register a new batch process.
  *
  * @param  array $args Arguments for the batch process.
- * @throws Exception Only post & user are accepted $args['type'].
+ * @throws Exception If the batch processor doesn't extend the Batch abstract class.
  */
 function register_batch_process( $args ) {
 	if ( empty( $args['type'] ) ) {
 		$args['type'] = '';
 	}
 
-	switch ( $args['type'] ) {
+	$batch_processor = get_default_batch_processor_for_type( $args['type'] );
+
+	/**
+	 * Filter the batch processor to be used with the batch process being registered.
+	 *
+	 * @param Batch  $batch_processor The batch processor to use, defaults to null.
+	 * @param string $type            Type of data for this batch process.
+	 * @param array  $args            Arguments for the batch process.
+	 */
+	$batch_processor = apply_filters( 'loco_register_batch_process_processor', $batch_processor, $args['type'], $args );
+
+	if ( empty( $batch_processor ) ) {
+		throw new Exception( sprintf(
+			__( 'Batch processor not found for type "%1$s"', 'locomotive' ),
+			$args['type']
+		) );
+	}
+
+	if ( ! is_subclass_of( $batch_processor, 'Rkv\Locomotive\Abstracts\Batch' ) ) {
+		throw new Exception( __( 'Batch processor must extend the Batch abstract class.', 'locomotive' ) );
+	}
+
+	$batch_processor->register( $args );
+}
+
+/**
+ * Returns the default batch processor used for a specific type of data.
+ *
+ * @param string $type Type of data to get a batch processor for.
+ *
+ * @return Batch|null The batch processor to use for the specified type, null if not a default type.
+ */
+function get_default_batch_processor_for_type( $type ) {
+	switch ( $type ) {
 		case 'post':
-			$batch_process = new Posts();
-			$batch_process->register( $args );
+			return new Posts();
 			break;
 
 		case 'user':
-			$batch_process = new Users();
-			$batch_process->register( $args );
+			return new Users();
 			break;
+
 		case 'site':
 			if ( is_multisite() ) {
-				$batch_process = new Sites();
-				$batch_process->register( $args );
+				return new Sites();
 			}
 			break;
 
 		case 'term':
-			$batch_process = new Terms();
-			$batch_process->register( $args );
+			return new Terms();
 			break;
 
 		case 'comment':
-			$batch_process = new Comments();
-			$batch_process->register( $args );
+			return new Comments();
 			break;
 	}
+
+	return null;
 }
 
 /**

tests/test-functions.php

@@ -3,6 +3,8 @@
 namespace Rkv\Locomotive\Tests;
 
 use WP_UnitTestCase;
+use Exception;
+use Rkv\Locomotive\Abstracts\Batch;
 use Rkv\Locomotive\Batches\Posts;
 use Rkv\Locomotive\Loader;
 
@@ -40,6 +42,67 @@ public function test_successful_register_batch() {
 		$this->assertCount( 2, $all_batches );
 	}
 
+	/**
+	 * Test that an Exception is thrown if a Batch Processor isn't found for a specified type.
+	 */
+	public function test_register_batch_process_throws_exception_if_no_batch_processor_for_type() {
+		$this->setExpectedException( Exception::class );
+
+		register_batch_process( array(
+			'name'     => 'My Failed Test Batch process',
+			'type'     => 'foo',
+			'callback' => __NAMESPACE__ . '\\my_callback_function',
+		) );
+	}
+
+	/**
+	 * Test that the batch processor filter works properly.
+	 */
+	public function test_batch_processor_filter_works() {
+		// Create a mock Batch object, set the expectation that register will be called
+		$mock_batch_processor = $this->getMockBuilder( Batch::class )->getMock();
+		$mock_batch_processor->expects( $this->once() )
+			->method( 'register' );
+
+		add_filter( 'loco_register_batch_process_processor', function ( $batch_processor, $type ) use ( $mock_batch_processor ) {
+			if ( 'foo' !== $type ) {
+				return $batch_processor;
+			}
+
+			return $mock_batch_processor;
+		}, 10, 2 );
+
+		register_batch_process( array(
+			'name'     => 'My Test Batch process',
+			'type'     => 'foo',
+			'callback' => __NAMESPACE__ . '\\my_callback_function',
+		) );
+	}
+
+	/**
+	 * Test that an Exception is thrown if a non-Batch Batch Processor is set in the filter.
+	 */
+	public function test_register_batch_process_throws_exception_if_batch_processor_is_not_a_batch() {
+		// Create a mock Batch object, set the expectation that register will be called
+		$mock_batch_processor = $this->getMockBuilder( \stdClass::class )->getMock();
+
+		add_filter( 'loco_register_batch_process_processor', function ( $batch_processor, $type ) use ( $mock_batch_processor ) {
+			if ( 'foo' !== $type ) {
+				return $batch_processor;
+			}
+
+			return $mock_batch_processor;
+		}, 10, 2 );
+
+		$this->setExpectedException( Exception::class );
+
+		register_batch_process( array(
+			'name'     => 'My Test Batch process',
+			'type'     => 'foo',
+			'callback' => __NAMESPACE__ . '\\my_callback_function',
+		) );
+	}
+
 	/**
 	 * Test that we can get all batches with proper data.
 	 */
@@ -127,6 +190,13 @@ public function test_elapsed_time() {
 		$this->assertInternalType( 'string', locomotive_time_ago( 1472852621 ) );
 	}
 
+	/**
+	 * Test the get_default_batch_processor_for_type function returns null if a non default type is provided.
+	 */
+	public function test_get_default_batch_processor_for_type_returns_null_for_non_default_type() {
+		$this->assertNull( get_default_batch_processor_for_type( 'foo' ) );
+	}
+
 	/**
 	 * Helper function to register a successful batch.
 	 *