Fix total count

assets/src/js/batch.jsx

@@ -74,7 +74,8 @@ var App = React.createClass( {
 		}
 
 		var self = this,
-			batchSlug = self.state.processing.batch;
+			batchSlug = self.state.processing.batch,
+			totalResults = self.state.processing.remote_data.total_num_results;
 
 		// If we open the modal and it was previously complete, clear it.
 		if ( 100 === this.state.processing.remote_data.progress ) {
@@ -86,13 +87,16 @@ var App = React.createClass( {
 			this.setState( { processing: this.state.processing, errors: [] } );
 		}
 
+
+
 		$.ajax( {
 			type: 'POST',
 			url: batch.ajaxurl,
 			data: {
 				batch_process: batchSlug,
 				nonce:         batch.nonce,
 				step:          currentStep,
+				total_num_results: totalResults,
 				action:        'run_batch',
 			},
 			dataType: 'json',
@@ -126,7 +130,7 @@ var App = React.createClass( {
 				// Determine if we have to run another step in the batch. Checks if there are more steps
 				// that need to run and makes sure the 'status' from the server is still 'running'.
 				if ( response.currentStep !== response.total_steps && 'running' === response.status.toLowerCase() ) {
-					self.runBatch( currentStep + 1 );
+					self.runBatch( response.current_step + 1 );
 				}
 			}
 		} ).fail( function () {

includes/abstracts/abstract-batch.php

@@ -87,6 +87,13 @@ abstract class Batch {
 	 */
 	public $total_num_results;
 
+	/**
+	 * Holds difference between total from client and total from query, if one exists.
+	 *
+	 * @var int
+	 */
+	public $difference_in_result_totals = 0;
+
 	/**
 	 * Errors from results
 	 *
@@ -136,8 +143,47 @@ abstract public function update_result_item_status( $result, $status );
 	 */
 	public function get_results() {
 		$this->args = wp_parse_args( $this->args, $this->default_args );
+		$results = $this->batch_get_results();
 		$this->calculate_offset();
-		return $this->batch_get_results();
+		return $results;
+	}
+
+	/**
+	 * Set the total number of results
+	 *
+	 * Uses a number passed from the client to the server and compares it to the total objects
+	 * pulled by the latest query. If the dataset is larger, we increase the total_num_results number.
+	 * Otherwise, keep it at the original (to acount for deletion / changes).
+	 *
+	 * @param int $total_from_query Total number of results from latest query.
+	 */
+	public function set_total_num_results( $total_from_query ) {
+		// If this is past step 1, the client is passing back the total number of results.
+		// This accounts for deletion / descructive actions to the data.
+		$total_from_request = isset( $_POST['total_num_results'] ) ? absint( $_POST['total_num_results'] ) : 0; // Input var okay.
+
+		// We need to check to see if there is any new data that has been added.
+		if ( $total_from_query > $total_from_request ) {
+			$this->total_num_results = (int) $total_from_query;
+		} else {
+			$this->total_num_results = (int) $total_from_request;
+		}
+
+		$this->record_change_if_totals_differ( $total_from_request, $total_from_query );
+	}
+
+	/**
+	 * If the amount of total records has changed, the amount is recorded so that it can
+	 * be applied to the offeset when it is calculated. This ensures that the offset takes into
+	 * account if new objects have been added or removed from the query.
+	 *
+	 * @param  int $total_from_request    Total number of results passed up from client.
+	 * @param  int $total_from_query      Total number of results retreived from query.
+	 */
+	public function record_change_if_totals_differ( $total_from_request, $total_from_query ) {
+		if ( $total_from_query !== $total_from_request && $total_from_request > 0 ) {
+			$this->difference_in_result_totals = $total_from_request - $total_from_query;
+		}
 	}
 
 	/**
@@ -146,7 +192,8 @@ public function get_results() {
 	public function calculate_offset() {
 		if ( 1 !== $this->current_step ) {
 			// Example: step 2: 1 * 10 = offset of 10, step 3: 2 * 10 = offset of 20.
-			$this->args['offset'] = ( ( $this->current_step - 1 ) * $this->args[ $this->per_batch_param ] );
+			// Also subtracting by results changed to account for deleted and added objects.
+			$this->args['offset'] = ( ( $this->current_step - 1 ) * $this->args[ $this->per_batch_param ] ) - $this->difference_in_result_totals;
 		}
 	}
 
@@ -271,8 +318,20 @@ public function run( $current_step ) {
 		$per_page = apply_filters( 'loco_batch_' . $this->slug . '_per_page', $per_page );
 
 		$total_steps = ceil( $this->total_num_results / $per_page );
+
 		if ( (int) $this->current_step === (int) $total_steps ) {
-			$this->update_status( 'finished' );
+
+			// Need to really check to make sure there were no results added while processing.
+			// In the case of destructive actions (i.e. deletion) there will be a gap equal to the per_page param.
+			// In all other cases, the difference in totals should equal total number of results.
+			// If neither of these are true, we need to run the last step over again.
+			$difference = $this->total_num_results - $this->difference_in_result_totals;
+			if ( $difference <= $per_page || $difference === $this->total_num_results ) {
+				$this->update_status( 'finished' );
+			} else {
+				$this->current_step = $this->current_step - 1;
+				$this->update_status( 'running' );
+			}
 		} else {
 			$this->update_status( 'running' );
 		}

includes/batches/class-batch-comments.php

@@ -39,7 +39,8 @@ class Comments extends Batch {
 	 */
 	public function batch_get_results() {
 		$query = new WP_Comment_Query( $this->args );
-		$this->total_num_results = $query->found_comments;
+		$total_comments = $query->found_comments;
+		$this->set_total_num_results( $total_comments );
 		return $query->get_comments();
 	}
 

includes/batches/class-batch-posts.php

@@ -39,7 +39,8 @@ class Posts extends Batch {
 	 */
 	public function batch_get_results() {
 		$query = new WP_Query( $this->args );
-		$this->total_num_results = $query->found_posts;
+		$total_posts = $query->found_posts;
+		$this->set_total_num_results( $total_posts );
 		return $query->get_posts();
 	}
 

includes/batches/class-batch-sites.php

@@ -39,7 +39,8 @@ class Sites extends Batch {
 	 */
 	public function batch_get_results() {
 		$query = new WP_Site_Query( $this->args );
-		$this->total_num_results = $query->found_sites;
+		$total_sites = $query->found_sites;
+		$this->set_total_num_results( $total_sites );
 		return $query->get_sites();
 	}
 

includes/batches/class-batch-terms.php

@@ -38,12 +38,10 @@ class Terms extends Batch {
 	 */
 	public function batch_get_results() {
 		$query = new WP_Term_Query( $this->args );
-
 		// Need to do a count query in order to get all possible terms as an integer.
 		$count_args = wp_parse_args( array( 'fields' => 'count', 'offset' => 0 ), $this->args );
 		$count = new WP_Term_Query( $count_args );
-		$this->total_num_results = $count->get_terms();
-
+		$this->set_total_num_results( $count->get_terms() );
 		return $query->get_terms();
 	}
 

includes/batches/class-batch-users.php

@@ -38,7 +38,8 @@ class Users extends Batch {
 	 */
 	public function batch_get_results() {
 		$query = new WP_User_Query( $this->args );
-		$this->total_num_results = $query->get_total();
+		$total_users = $query->get_total();
+		$this->set_total_num_results( $total_users );
 		return $query->get_results();
 	}
 

tests/types/test-batch-posts.php

@@ -122,6 +122,37 @@ public function test_offset_run() {
 		$this->assertEquals( 'finished', $batch_status['status'] );
 	}
 
+	/**
+	 * Test that batch gets run.
+	 */
+	public function test_run_with_destructive_callback() {
+
+		$post_batch = new Posts();
+
+		$post_batch->register( array(
+			'name'     => 'Hey there',
+			'type'     => 'post',
+			'callback' => __NAMESPACE__ . '\\my_callback_delete_post',
+			'args'     => array(
+				'posts_per_page' => 5,
+				'post_type'      => 'post',
+			),
+		) );
+
+		// Simulate running twice with pass back of results number from client.
+		$post_batch->run( 1 );
+		$_POST['total_num_results'] = 10;
+		$post_batch->run( 2 );
+
+		// Check that all posts have been deleted.
+		$all_posts = get_posts( array( 'post_type' => 'post', 'post_status' => 'publish', 'posts_per_page' => -1 ) );
+		$this->assertCount( 0, $all_posts );
+
+		// Ensure that we are still getting a finished message.
+		$batch_status = get_option( 'loco_batch_' . $post_batch->slug );
+		$this->assertEquals( 'finished', $batch_status['status'] );
+	}
+
 }
 
 /**
@@ -132,3 +163,12 @@ public function test_offset_run() {
 function my_post_callback_function_test( $result ) {
 	update_post_meta( $result->ID, 'custom-key', 'my-value' );
 }
+
+/**
+ * Callback function with destructive action (deletion).
+ *
+ * @param  WP_Post $result Result item.
+ */
+function my_callback_delete_post( $result ) {
+	wp_delete_post( $result->ID, true );
+}

tests/types/test-batch-terms.php

@@ -110,7 +110,7 @@ public function test_terms_offset_run() {
 			'type'     => 'term',
 			'callback' => __NAMESPACE__ . '\\my_term_callback_function_test',
 			'args'     => array(
-				'number' => 5,
+				'number' => 3,
 				'offset' => 5,
 				'taxonomy' => 'category',
 				'hide_empty' => false,

tests/types/test-batch-users.php

@@ -19,7 +19,7 @@ public function tearDown() {
 	 */
 	public function setUp() {
 		parent::setUp();
-		$this->users = $this->factory->user->create_many( 5 );
+		$this->users = $this->factory->user->create_many( 9 );
 	}
 
 
@@ -70,7 +70,7 @@ public function test_clear_user_result_status() {
 			'type'     => 'user',
 			'callback' => __NAMESPACE__ . '\\my_user_callback_function_test',
 			'args'     => array(
-				'number' => 7,
+				'number' => 10,
 			),
 		) );