From 377ac507735d369efed974adb706d86591aabd32 Mon Sep 17 00:00:00 2001
From: Mateusz Uzdowski <mateusz@silverstripe.com>
Date: Thu, 7 Jun 2012 10:26:27 +1200
Subject: [PATCH] ENHNANCEMENT: allow overriding of PJAX fragments included in
 the response.

---
 control/PjaxResponseNegotiator.php           | 46 +++++++++++++++-----
 tests/control/PjaxResponseNegotiatorTest.php | 18 +++++++-
 2 files changed, 53 insertions(+), 11 deletions(-)

diff --git a/control/PjaxResponseNegotiator.php b/control/PjaxResponseNegotiator.php
index 2aabeaef4..6f24a1ddc 100644
--- a/control/PjaxResponseNegotiator.php
+++ b/control/PjaxResponseNegotiator.php
@@ -22,6 +22,11 @@ class PjaxResponseNegotiator {
 
 	protected $response = null;
 
+	/**
+	 * Overriden fragments (if any). Otherwise uses fragments from the request.
+	 */
+	protected $fragmentOverride = null;
+
 	/**
 	 * @param RequestHandler $controller
 	 * @param SS_HTTPResponse An existing response to reuse (optional)
@@ -49,6 +54,7 @@ class PjaxResponseNegotiator {
 	 * @param array $extraCallbacks List of anonymous functions or callables returning either a string
 	 * or SS_HTTPResponse, keyed by their fragment identifier. The 'default' key can
 	 * be used as a fallback for non-ajax responses.
+	 * @param array $fragmentOverride Change the response fragments.
 	 * @return SS_HTTPResponse
 	 */
 	public function respond(SS_HTTPRequest $request, $extraCallbacks = array()) {
@@ -57,21 +63,28 @@ class PjaxResponseNegotiator {
 		$response = $this->getResponse();
 		
 		$responseParts = array();
-		if($fragmentStr = $request->getHeader('X-Pjax')) {
+
+		if (isset($this->fragmentOverride)) {
+			$fragments = $this->fragmentOverride;
+		} elseif ($fragmentStr = $request->getHeader('X-Pjax')) {
 			$fragments = explode(',', $fragmentStr);
-			foreach($fragments as $fragment) {
-				if(isset($callbacks[$fragment])) {
-					$responseParts[$fragment] = call_user_func($callbacks[$fragment]);
-				} else {
-					throw new SS_HTTPResponse_Exception("X-Pjax = '$fragment' not supported for this URL.", 400);
-				}
-			}
-			$response->setBody(Convert::raw2json($responseParts));
-			$response->addHeader('Content-Type', 'text/json');
 		} else {
 			if($request->isAjax()) throw new SS_HTTPResponse_Exception("Ajax requests to this URL require an X-Pjax header.", 400);
 			$response->setBody(call_user_func($callbacks['default']));
+			return $response;
 		}
+
+		// Execute the fragment callbacks and build the response.
+		foreach($fragments as $fragment) {
+			if(isset($callbacks[$fragment])) {
+				$responseParts[$fragment] = call_user_func($callbacks[$fragment]);
+			} else {
+				throw new SS_HTTPResponse_Exception("X-Pjax = '$fragment' not supported for this URL.", 400);
+			}
+		}
+		$response->setBody(Convert::raw2json($responseParts));
+		$response->addHeader('Content-Type', 'text/json');
+
 		return $response;
 	}
 
@@ -82,4 +95,17 @@ class PjaxResponseNegotiator {
 	public function setCallback($fragment, $callback) {
 		$this->callbacks[$fragment] = $callback;
 	}
+
+	/**
+	 * Set up fragment overriding - will completely replace the incoming fragments.
+	 *
+	 * @param array $fragments Fragments to insert.
+	 */
+	public function overrideFragments($fragments) {
+		if (!is_array($fragments)) throw new InvalidArgumentException();
+
+		$this->fragmentOverride = $fragments;
+
+		return $this;
+	}
 }
diff --git a/tests/control/PjaxResponseNegotiatorTest.php b/tests/control/PjaxResponseNegotiatorTest.php
index c1162b622..be697b771 100644
--- a/tests/control/PjaxResponseNegotiatorTest.php
+++ b/tests/control/PjaxResponseNegotiatorTest.php
@@ -38,4 +38,20 @@ class PjaxResponseNegotiatorTest extends SapphireTest {
 		$this->assertEquals('otherfragment response', $json->otherfragment);
 	}
 
-}
\ No newline at end of file
+	function testFragmentsOverride() {
+		$negotiator = new PjaxResponseNegotiator(array(
+			'alpha' => function() {return 'alpha response';},
+			'beta' => function() {return 'beta response';}
+		));
+
+		$request = new SS_HTTPRequest('GET', '/');
+		$request->addHeader('X-Pjax', 'alpha');
+		$request->addHeader('Accept', 'text/json');
+
+		$response = $negotiator->overrideFragments(array('beta'))->respond($request);
+		$json = json_decode( $response->getBody());
+		$this->assertFalse(isset($json->alpha));
+		$this->assertObjectHasAttribute('beta', $json);
+	}
+
+}