FIX Get page ID from document set and include it in $NewLink for add document buttons

This helps the module to determine where to return a user to after they've uploaded a file, e.g. either back to a document set in a page context, back to a document set in the ModelAdmin context or back to the ModelAdmin for documents with no context.

code/cms/DMSDocumentAddController.php

@@ -35,26 +35,21 @@ class DMSDocumentAddController extends LeftAndMain
     public function currentPage()
     {
         $id = $this->currentPageID();
-
+        if ($id === 0) {
-        if ($id && is_numeric($id) && $id > 0) {
+            return SiteTree::singleton();
-            return Versioned::get_by_stage('SiteTree', 'Stage', sprintf(
-                'ID = %s',
-                (int) $id
-            ))->first();
-        } else {
-            // ID is either '0' or 'root'
-            return singleton('SiteTree');
         }
+        return parent::currentPage();
     }
 
     /**
-     * Return fake-ID "root" if no ID is found (needed to upload files into the root-folder)
+     * Return fake-ID "root" if no ID is found (needed to upload files into the root-folder). Otherwise the page ID
+     * is passed in from the {@link DMSGridFieldAddNewButton}.
      *
      * @return int
      */
     public function currentPageID()
     {
-        return ($result = parent::currentPageID()) === null ? 0 : $result;
+        return (int) $this->getRequest()->getVar('page_id');
     }
 
     /**
@@ -177,11 +172,11 @@ class DMSDocumentAddController extends LeftAndMain
     }
 
     /**
-     * Returns the link to be used to return the user after uploading a document. If a document set ID (dsid) is present
+     * Returns the link to be used to return the user after uploading a document. Scenarios:
-     * then it will be redirected back to that document set in a page.
      *
-     * If no document set ID is present then we assume that it has been added from the model admin, so redirect back to
+     * 1) Page context: page ID and document set ID provided, redirect back to the page and document set
-     * that instead.
+     * 2) Document set context: no page ID, document set ID provided, redirect back to document set in ModelAdmin
+     * 3) Document context: no page ID and no document set ID provided, redirect back to documents in ModelAdmin
      *
      * @return string
      */

code/cms/DMSGridFieldAddNewButton.php

@@ -32,6 +32,14 @@ class DMSGridFieldAddNewButton extends GridFieldAddNewButton implements GridFiel
         $link = singleton('DMSDocumentAddController')->Link();
         if ($this->getDocumentSetId()) {
             $link = Controller::join_links($link, '?dsid=' . $this->getDocumentSetId());
+
+            // Look for an associated page, but only share it if we're editing in a page context
+            $set = DMSDocumentSet::get()->byId($this->getDocumentSetId());
+            if ($set && $set->exists() && $set->Page()->exists()
+                && Controller::curr() instanceof CMSPageEditController
+            ) {
+                $link = Controller::join_links($link, '?page_id=' . $set->Page()->ID);
+            }
         }
 
         $data = new ArrayData(array(

tests/DMSGridFieldAddNewButtonTest.php

@@ -1,47 +0,0 @@
-<?php
-
-class DMSGridFieldAddNewButtonTest extends SapphireTest
-{
-    protected static $fixture_file = 'dmstest.yml';
-
-    /**
-     * @var DMSGridFieldAddNewButton
-     */
-    protected $button;
-
-    /**
-     * @var GridField
-     */
-    protected $gridField;
-
-    public function setUp()
-    {
-        parent::setUp();
-
-        $fakeList = DMSDocument::get();
-        $this->gridField = GridField::create('TestGridField', false, $fakeList);
-        $this->button = new DMSGridFieldAddNewButton;
-    }
-
-    /**
-     * Test that when no document set ID is present then it is not added to the URL for "add document"
-     */
-    public function testNoPageIdInAddUrlWhenNotProvided()
-    {
-        $fragments = $this->button->getHTMLFragments($this->gridField);
-        $result = array_pop($fragments)->getValue();
-        $this->assertNotContains('?dsid', $result);
-    }
-
-    /**
-     * Test that when a document set ID is provided, it is added onto the "add document" link
-     */
-    public function testPageIdAddedToLinkWhenProvided()
-    {
-        $this->button->setDocumentSetId(123);
-
-        $fragments = $this->button->getHTMLFragments($this->gridField);
-        $result = array_pop($fragments)->getValue();
-        $this->assertContains('?dsid=123', $result);
-    }
-}

tests/cms/DMSDocumentAddControllerTest.php

@@ -27,7 +27,7 @@ class DMSDocumentAddControllerTest extends FunctionalTest
 
         $this->assertInstanceOf('SiteTree', $this->controller->currentPage());
         $this->assertEmpty($this->controller->currentPage()->ID);
-        $this->controller->setRequest(new SS_HTTPRequest('GET', '/', array('ID' => $page->ID)));
+        $this->controller->setRequest(new SS_HTTPRequest('GET', '/', array('page_id' => $page->ID)));
         $this->assertEquals($page->ID, $this->controller->currentPage()->ID, 'Specified page is loaded and returned');
     }
 
@@ -74,7 +74,7 @@ class DMSDocumentAddControllerTest extends FunctionalTest
         $this->assertContains('123', $this->controller->Backlink());
 
         // Has page ID and document set ID
-        $request = new SS_HTTPRequest('GET', '/', array('dsid' => 123, 'ID' => 234));
+        $request = new SS_HTTPRequest('GET', '/', array('dsid' => 123, 'page_id' => 234));
         $this->controller->setRequest($request);
         $this->assertContains('admin/pages', $this->controller->Backlink());
         $this->assertContains('123', $this->controller->Backlink());

tests/cms/DMSGridFieldAddNewButtonTest.php

@@ -0,0 +1,79 @@
+<?php
+
+class DMSGridFieldAddNewButtonTest extends SapphireTest
+{
+    protected static $fixture_file = 'dms/tests/dmstest.yml';
+
+    /**
+     * @var DMSGridFieldAddNewButton
+     */
+    protected $button;
+
+    /**
+     * @var GridField
+     */
+    protected $gridField;
+
+    public function setUp()
+    {
+        parent::setUp();
+
+        $fakeList = DMSDocument::get();
+        $this->gridField = GridField::create('TestGridField', false, $fakeList);
+        $this->button = new DMSGridFieldAddNewButton;
+    }
+
+    /**
+     * Test that when no document set ID is present then it is not added to the URL for "add document"
+     */
+    public function testNoDocumentSetIdInAddUrlWhenNotProvided()
+    {
+        $this->assertNotContains('?dsid', $this->getButtonHtml());
+    }
+
+    /**
+     * Test that when a document set ID is provided, it is added onto the "add document" link
+     */
+    public function testDocumentSetIdAddedToLinkWhenProvided()
+    {
+        $this->button->setDocumentSetId(123);
+        $this->assertContains('?dsid=123', $this->getButtonHtml());
+    }
+
+    /**
+     * If a set is saved and associated to a page, that page's ID should be added to the "add document" link to help
+     * to ensure the user gets redirected back to the correct place afterwards
+     */
+    public function testPageIdIsAddedWhenAvailableViaDocumentSetRelationship()
+    {
+        $set = $this->objFromFixture('DMSDocumentSet', 'ds1');
+        $this->button->setDocumentSetId($set->ID);
+
+        $controller = new ContentController;
+        $controller->pushCurrent();
+
+        $result = $this->getButtonHtml();
+        $this->assertContains('dsid=' . $set->ID, $result, 'Add new button contains document set ID');
+        $this->assertNotContains('page_id=' . $set->Page()->ID, $result, 'No page ID when not editing in page context');
+
+        $controller = new CMSPageEditController;
+        $controller->pushCurrent();
+
+        $this->assertContains(
+            'page_id=' . $set->Page()->ID,
+            $this->getButtonHtml(),
+            'Button contains page ID when in edit page context'
+        );
+    }
+
+    /**
+     * Returns the HTML result of the "add new" button, used for DRY in test class
+     *
+     * @return string
+     */
+    protected function getButtonHtml()
+    {
+        $fragments = $this->button->getHTMLFragments($this->gridField);
+        return array_pop($fragments)->getValue();
+    }
+}