From 5a75a4ef49aa56718d7eb6a828eb7164f28460a7 Mon Sep 17 00:00:00 2001
From: Milen Petrinski <gonzo@greatgonzo.net>
Date: Sat, 20 Mar 2021 18:26:02 +0200
Subject: [PATCH 1/2] Restore current blog when the_post triggers outside the
 loop.

---
 bin/install-wp-tests.sh                         |  0
 .../classes/Indexable/Post/QueryIntegration.php | 17 ++++++++++++-----
 2 files changed, 12 insertions(+), 5 deletions(-)
 mode change 100644 => 100755 bin/install-wp-tests.sh

diff --git a/bin/install-wp-tests.sh b/bin/install-wp-tests.sh
old mode 100644
new mode 100755
diff --git a/includes/classes/Indexable/Post/QueryIntegration.php b/includes/classes/Indexable/Post/QueryIntegration.php
index af999b0bbf..7621ecadd1 100644
--- a/includes/classes/Indexable/Post/QueryIntegration.php
+++ b/includes/classes/Indexable/Post/QueryIntegration.php
@@ -51,7 +51,7 @@ public function __construct() {
 		add_action( 'loop_end', array( $this, 'maybe_restore_blog' ), 10, 1 );
 
 		// Properly switch to blog if necessary
-		add_action( 'the_post', array( $this, 'maybe_switch_to_blog' ), 10, 1 );
+		add_action( 'the_post', array( $this, 'maybe_switch_to_blog' ), 10, 2 );
 
 		// Sets the correct value for found_posts
 		add_filter( 'found_posts', array( $this, 'found_posts' ), 10, 2 );
@@ -133,10 +133,11 @@ public function get_switched() {
 	/**
 	 * Switch to the correct site if the post site id is different than the actual one
 	 *
-	 * @param WP_Post $post Post object
+	 * @param WP_Post  $post Post object
+	 * @param WP_Query $query WP_Query instance
 	 * @since 0.9
 	 */
-	public function maybe_switch_to_blog( $post ) {
+	public function maybe_switch_to_blog( $post, $query ) {
 		if ( ! is_multisite() ) {
 			// @codeCoverageIgnoreStart
 			return;
@@ -154,9 +155,15 @@ public function maybe_switch_to_blog( $post ) {
 
 			$this->switched = $post->site_id;
 
-			remove_action( 'the_post', array( $this, 'maybe_switch_to_blog' ), 10, 1 );
+			remove_action( 'the_post', array( $this, 'maybe_switch_to_blog' ), 10, 2 );
 			setup_postdata( $post );
-			add_action( 'the_post', array( $this, 'maybe_switch_to_blog' ), 10, 1 );
+			add_action( 'the_post', array( $this, 'maybe_switch_to_blog' ), 10, 2 );
+
+			if ( $this->switched && ! $query->in_the_loop ) {
+				restore_current_blog();
+
+				$this->switched = false;
+			}
 		}
 
 	}

From a4425468c84745c8b1692d3bba68f0ee5a72e276 Mon Sep 17 00:00:00 2001
From: Milen Petrinski <gonzo@greatgonzo.net>
Date: Mon, 22 Mar 2021 08:50:51 +0200
Subject: [PATCH 2/2] Fix post indexable tests, test the new behaviour of
 maybe_switch_to_blog.

---
 tests/php/indexables/TestPost.php | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/tests/php/indexables/TestPost.php b/tests/php/indexables/TestPost.php
index 4e5aa48974..94e014a964 100644
--- a/tests/php/indexables/TestPost.php
+++ b/tests/php/indexables/TestPost.php
@@ -5663,14 +5663,22 @@ public function testMaybeSwitchToBlog() {
 		// test the function. Try accessing the 2nd post from the 1st blog.
 		$query_integration = new \ElasticPress\Indexable\Post\QueryIntegration();
 
+		// This should not switch to the 2nd site because the query is not in the loop.
+		$query_integration->maybe_switch_to_blog( $blog_2_post, $query );
+
+		$this->assertFalse( $query_integration->get_switched() );
+
+		// To switch sites the query must be in the loop.
+		$query->in_the_loop = true;
+
 		// This should switch to the 2nd site.
-		$query_integration->maybe_switch_to_blog( $blog_2_post );
+		$query_integration->maybe_switch_to_blog( $blog_2_post, $query );
 
 		$this->assertSame( $blog_2_post->site_id, $query_integration->get_switched() );
 
 		// Now we're in "switched" mode, try getting the post from the
 		// 1st site, should switch back.
-		$query_integration->maybe_switch_to_blog( $blog_1_post );
+		$query_integration->maybe_switch_to_blog( $blog_1_post, $query );
 
 		$this->assertSame( $blog_1_post->site_id, $query_integration->get_switched() );