From cfd263f653d20fab0c30ebeb29c385b999d33745 Mon Sep 17 00:00:00 2001
From: Jacob Bennett <Jacob.J.Bennett@gmail.com>
Date: Wed, 7 Jun 2017 14:45:59 -0500
Subject: [PATCH] Limit to Links

* Limit to links

* Added test case

* int type declaration removed

* int type declaration removed

* test fix

* extra condition removed

* typo

* modify tests
---
 src/Middleware/AddHttp2ServerPush.php | 10 ++++++----
 tests/AddHttp2ServerPushTest.php      | 11 +++++++++++
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/src/Middleware/AddHttp2ServerPush.php b/src/Middleware/AddHttp2ServerPush.php
index ac3fd26..108e599 100644
--- a/src/Middleware/AddHttp2ServerPush.php
+++ b/src/Middleware/AddHttp2ServerPush.php
@@ -25,7 +25,7 @@ class AddHttp2ServerPush
      *
      * @return mixed
      */
-    public function handle(Request $request, Closure $next)
+    public function handle(Request $request, Closure $next, $limit = null)
     {
         $response = $next($request);
 
@@ -33,7 +33,7 @@ public function handle(Request $request, Closure $next)
             return $response;
         }
 
-        $this->generateAndAttachLinkHeaders($response);
+        $this->generateAndAttachLinkHeaders($response, $limit);
 
         return $response;
     }
@@ -43,13 +43,15 @@ public function handle(Request $request, Closure $next)
      *
      * @return $this
      */
-    protected function generateAndAttachLinkHeaders(Response $response)
+    protected function generateAndAttachLinkHeaders(Response $response, $limit = null)
     {
         $headers = $this->fetchLinkableNodes($response)
             ->flatten(1)
             ->map(function ($url) {
                 return $this->buildLinkHeaderString($url);
-            })->filter()
+            })
+            ->filter()
+            ->take($limit)
             ->implode(',');
 
         if (!empty(trim($headers))) {
diff --git a/tests/AddHttp2ServerPushTest.php b/tests/AddHttp2ServerPushTest.php
index c15eed1..6e6e912 100644
--- a/tests/AddHttp2ServerPushTest.php
+++ b/tests/AddHttp2ServerPushTest.php
@@ -93,6 +93,17 @@ public function it_will_not_return_a_push_header_for_inline_js()
         $this->assertFalse($this->isServerPushResponse($response));
     }
 
+    /** @test */
+    public function it_will_return_limit_count_of_links()
+    {
+        $request = new Request();
+        $limit = 2;
+
+        $response = $this->middleware->handle($request, $this->getNext('pageWithImages'), $limit);
+
+        $this->assertCount($limit, explode(",", $response->headers->get('link')));
+    }
+
     /**
      * @param string $pageName
      *