From b3ce59ec563b16d00407669812f075dbcab7c4c3 Mon Sep 17 00:00:00 2001
From: Weston Ruter <westonruter@google.com>
Date: Mon, 24 Jun 2024 13:23:05 -0700
Subject: [PATCH] Try considering PHP_OUTPUT_HANDLER_FINAL in output buffer
 callback

---
 .../optimization-detective/optimization.php   | 23 +++++-----
 .../tests/test-optimization.php               | 42 ++++++++++++++++++-
 2 files changed, 54 insertions(+), 11 deletions(-)

diff --git a/plugins/optimization-detective/optimization.php b/plugins/optimization-detective/optimization.php
index 442a7a6226..2e75924826 100644
--- a/plugins/optimization-detective/optimization.php
+++ b/plugins/optimization-detective/optimization.php
@@ -32,16 +32,19 @@
  */
 function od_buffer_output( string $passthrough ): string {
 	ob_start(
-		static function ( string $output ): string {
-			/**
-			 * Filters the template output buffer prior to sending to the client.
-			 *
-			 * @since 0.1.0
-			 *
-			 * @param string $output Output buffer.
-			 * @return string Filtered output buffer.
-			 */
-			return (string) apply_filters( 'od_template_output_buffer', $output );
+		static function ( string $output, ?int $phase ): string {
+			if ( $phase & PHP_OUTPUT_HANDLER_FINAL ) {
+				/**
+				 * Filters the template output buffer prior to sending to the client.
+				 *
+				 * @since 0.1.0
+				 *
+				 * @param string $output Output buffer.
+				 * @return string Filtered output buffer.
+				 */
+				$output = (string) apply_filters( 'od_template_output_buffer', $output );
+			}
+			return $output;
 		}
 	);
 	return $passthrough;
diff --git a/plugins/optimization-detective/tests/test-optimization.php b/plugins/optimization-detective/tests/test-optimization.php
index ff62951a5c..af67949a05 100644
--- a/plugins/optimization-detective/tests/test-optimization.php
+++ b/plugins/optimization-detective/tests/test-optimization.php
@@ -54,10 +54,12 @@ public function test_od_buffer_output(): void {
 		// buffer callback. See <https://stackoverflow.com/a/61439514/93579>.
 		ob_start();
 
+		$filter_invoked = false;
 		add_filter(
 			'od_template_output_buffer',
-			function ( $buffer ) use ( $original, $expected ) {
+			function ( $buffer ) use ( $original, $expected, &$filter_invoked ) {
 				$this->assertSame( $original, $buffer );
+				$filter_invoked = true;
 				return $expected;
 			}
 		);
@@ -71,6 +73,44 @@ function ( $buffer ) use ( $original, $expected ) {
 
 		$buffer = ob_get_clean(); // Get the buffer from our wrapper output buffer.
 		$this->assertSame( $expected, $buffer );
+		$this->assertTrue( $filter_invoked );
+	}
+
+	/**
+	 * Make output is buffered and that it is also filtered.
+	 *
+	 * @covers ::od_buffer_output
+	 */
+	public function test_od_buffer_output_not_finalized(): void {
+		$original = 'Hello My World!';
+		$override = 'Ciao mondo!';
+
+		// In order to test, a wrapping output buffer is required because ob_get_clean() does not invoke the output
+		// buffer callback. See <https://stackoverflow.com/a/61439514/93579>.
+		ob_start();
+
+		$filter_invoked = false;
+		add_filter(
+			'od_template_output_buffer',
+			function ( $buffer ) use ( $original, &$filter_invoked ) {
+				$this->assertSame( $original, $buffer );
+				$filter_invoked = true;
+				return '¡Hola Mi Mundo!';
+			}
+		);
+
+		$original_ob_level = ob_get_level();
+		od_buffer_output( '' );
+		$this->assertSame( $original_ob_level + 1, ob_get_level(), 'Expected call to ob_start().' );
+		echo $original;
+
+		ob_clean(); // Note the lack of flush here.
+		echo $override;
+
+		$buffer = ob_get_clean(); // Get the buffer from our wrapper output buffer.
+
+		$this->assertFalse( $filter_invoked );
+		$this->assertSame( $override, $buffer );
 	}
 
 	/**