Force is_amp_endpoint to true when rendering.

If amp_render_post is called directly outside of the standard endpoint, is_amp_endpoint will return false, which is not ideal for any code that expects to run in an AMP context (e.g. inside shortcodes). Let's force the value to be true while we render AMP content.
jritterbush · Dec 14, 2017 · eb1d393 · eb1d393
1 parent d72dca3
commit eb1d393
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 1 deletion.
diff --git a/amp.php b/amp.php
@@ -192,6 +192,12 @@ function amp_render_post( $post ) {
 	}
 	$post_id = $post->ID;
 
+	// If amp_render_post is called directly outside of the standard endpoint, is_amp_endpoint will return false, which is not ideal for any code that expects to run in an AMP context.
+	// Let's force the value to be true while we render AMP.
+	add_filter( 'is_amp_endpoint', '__return_true' );
+
+	amp_load_classes();
+
 	/**
 	 * Fires before rendering a post in AMP.
 	 *
@@ -204,6 +210,8 @@ function amp_render_post( $post ) {
 	amp_add_post_template_actions();
 	$template = new AMP_Post_Template( $post );
 	$template->load();
+
+	remove_filter( 'is_amp_endpoint', '__return_true' );
 }
 
 /**

diff --git a/includes/amp-helper-functions.php b/includes/amp-helper-functions.php
@@ -78,7 +78,9 @@ function is_amp_endpoint() {
 		_doing_it_wrong( __FUNCTION__, sprintf( esc_html__( "is_amp_endpoint() was called before the 'parse_query' hook was called. This function will always return 'false' before the 'parse_query' hook is called.", 'amp' ) ), '0.4.2' );
 	}
 
-	return false !== get_query_var( AMP_QUERY_VAR, false );
+	$is_amp_endpoint = false !== get_query_var( AMP_QUERY_VAR, false );
+
+	return apply_filters( 'is_amp_endpoint', $is_amp_endpoint );
 }
 
 /**

diff --git a/tests/test-amp-render-post.php b/tests/test-amp-render-post.php
@@ -21,4 +21,33 @@ public function test__valid_post() {
 		$this->assertContains( '<html amp', $amp_rendered, 'Response does not include html tag with amp attribute.' );
 		$this->assertEquals( 1, did_action( 'pre_amp_render_post', 'pre_amp_render_post action fire either did not fire or fired too many times.' ) );
 	}
+
+	public function test__is_amp_endpoint() {
+		$user_id = $this->factory->user->create();
+		$post_id = $this->factory->post->create( array(
+			'post_author' => $user_id,
+		) );
+
+		$pre_is_amp_endpoint = is_amp_endpoint();
+
+		add_action( 'pre_amp_render_post', array( $this, '_check_is_amp_endpoint' ) );
+
+		// Need to use ob here since the method echos
+		ob_start();
+		amp_render_post( $post_id );
+		$amp_rendered = ob_get_clean();
+
+		$mid_is_amp_endpoint = $GLOBALS['mid_is_amp_endpoint'];
+		$post_is_amp_endpoint = is_amp_endpoint();
+
+		$this->assertFalse( $pre_is_amp_endpoint, 'is_amp_endpoint was not defaulting to false before amp_render_post' );
+		$this->assertTrue( $mid_is_amp_endpoint, 'is_amp_endpoint was not forced to true during amp_render_post' );
+		$this->assertFalse( $post_is_amp_endpoint, 'is_amp_endpoint was not reset after amp_render_post' );
+
+		unset( $GLOBALS['mid_is_amp_endpoint'] );
+	}
+
+	public function _check_is_amp_endpoint() {
+		$GLOBALS['mid_is_amp_endpoint'] = is_amp_endpoint();
+	}
 }