From deb0d68875dac62be21843934a0c146c9a54e383 Mon Sep 17 00:00:00 2001
From: n0nag0n <n0nag0n@sky-9.com>
Date: Thu, 11 Apr 2024 23:04:28 -0600
Subject: [PATCH] added plain stream() method

---
 flight/Engine.php                 | 13 ++++++++-----
 flight/net/Route.php              | 11 +++++++++++
 tests/FlightTest.php              | 24 ++++++++++++++++++++++++
 tests/server/LayoutMiddleware.php |  3 ++-
 tests/server/index.php            | 12 ++++++++++++
 5 files changed, 57 insertions(+), 6 deletions(-)

diff --git a/flight/Engine.php b/flight/Engine.php
index 0af8765..0d32f0a 100644
--- a/flight/Engine.php
+++ b/flight/Engine.php
@@ -487,13 +487,16 @@ class Engine
 
             // If this route is to be streamed, we need to output the headers now
             if ($route->is_streamed === true) {
-                $response->status($route->streamed_headers['status'] ?? 200);
-                unset($route->streamed_headers['status']);
+                if (count($route->streamed_headers) > 0) {
+                    $response->status($route->streamed_headers['status'] ?? 200);
+                    unset($route->streamed_headers['status']);
+                    foreach ($route->streamed_headers as $header => $value) {
+                        $response->header($header, $value);
+                    }
+                }
+
                 $response->header('X-Accel-Buffering', 'no');
                 $response->header('Connection', 'close');
-                foreach ($route->streamed_headers as $header => $value) {
-                    $response->header($header, $value);
-                }
 
                 // We obviously don't know the content length right now. This must be false.
                 $response->content_length = false;
diff --git a/flight/net/Route.php b/flight/net/Route.php
index 57ba2fc..4d005b9 100644
--- a/flight/net/Route.php
+++ b/flight/net/Route.php
@@ -238,6 +238,17 @@ class Route
         return $this;
     }
 
+    /**
+     * If the response should be streamed
+     *
+     * @return self
+     */
+    public function stream(): self
+    {
+        $this->is_streamed = true;
+        return $this;
+    }
+
     /**
      * This will allow the response for this route to be streamed.
      *
diff --git a/tests/FlightTest.php b/tests/FlightTest.php
index 042b6bb..d84336b 100644
--- a/tests/FlightTest.php
+++ b/tests/FlightTest.php
@@ -280,6 +280,30 @@ class FlightTest extends TestCase
     }
 
     public function testStreamRoute()
+    {
+        $response_mock = new class extends Response {
+            public function setRealHeader(string $header_string, bool $replace = true, int $response_code = 0): Response
+            {
+                return $this;
+            }
+        };
+        $mock_response_class_name = get_class($response_mock);
+        Flight::register('response', $mock_response_class_name);
+        Flight::route('/stream', function () {
+            echo 'stream';
+        })->stream();
+        Flight::request()->url = '/stream';
+        $this->expectOutputString('stream');
+        Flight::start();
+        $this->assertEquals('', Flight::response()->getBody());
+        $this->assertEquals([
+            'X-Accel-Buffering' => 'no',
+            'Connection' => 'close'
+        ], Flight::response()->getHeaders());
+        $this->assertEquals(200, Flight::response()->status());
+    }
+
+    public function testStreamRouteWithHeaders()
     {
         $response_mock = new class extends Response {
             public function setRealHeader(string $header_string, bool $replace = true, int $response_code = 0): Response
diff --git a/tests/server/LayoutMiddleware.php b/tests/server/LayoutMiddleware.php
index b89c4e0..d23e81c 100644
--- a/tests/server/LayoutMiddleware.php
+++ b/tests/server/LayoutMiddleware.php
@@ -76,7 +76,8 @@ class LayoutMiddleware
 <li><a href="/jsonp?jsonp=myjson">JSONP</a></li>
 <li><a href="/halt">Halt</a></li>
 <li><a href="/redirect">Redirect</a></li>
-<li><a href="/streamResponse">Stream</a></li>
+<li><a href="/streamResponse">Stream Plain</a></li>
+<li><a href="/streamWithHeaders">Stream Headers</a></li>
 <li><a href="/overwrite">Overwrite Body</a></li>
 <li><a href="/redirect/before%2Fafter">Slash in Param</a></li>
 <li><a href="/わたしはひとです">UTF8 URL</a></li>
diff --git a/tests/server/index.php b/tests/server/index.php
index d9d16bc..a983590 100644
--- a/tests/server/index.php
+++ b/tests/server/index.php
@@ -122,7 +122,19 @@ Flight::group('', function () {
             ob_flush();
         }
         echo "is successful!!";
+    })->stream();
+
+    // Test 12: Redirect with status code
+    Flight::route('/streamWithHeaders', function () {
+        echo "Streaming a response";
+        for ($i = 1; $i <= 50; $i++) {
+            echo ".";
+            usleep(50000);
+            ob_flush();
+        }
+        echo "is successful!!";
     })->streamWithHeaders(['Content-Type' => 'text/html', 'status' => 200 ]);
+
     // Test 14: Overwrite the body with a middleware
     Flight::route('/overwrite', function () {
         echo '<span id="infotext">Route text:</span> This route status is that it <span style="color:red; font-weight: bold;">failed</span>';