From 7f7fe73b772ebf652217c64d0767c66740f39709 Mon Sep 17 00:00:00 2001
From: Daniel Hensby <dan@hens.by>
Date: Mon, 25 Sep 2017 13:55:33 +0100
Subject: [PATCH] Reduce dependence on session state for accessing draft stages

---
 src/Dev/FunctionalTest.php | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/src/Dev/FunctionalTest.php b/src/Dev/FunctionalTest.php
index 7e14e944c..add41d955 100644
--- a/src/Dev/FunctionalTest.php
+++ b/src/Dev/FunctionalTest.php
@@ -4,6 +4,7 @@ namespace SilverStripe\Dev;
 
 use SilverStripe\Control\Controller;
 use SilverStripe\Control\Director;
+use SilverStripe\Control\HTTP;
 use SilverStripe\Control\Session;
 use SilverStripe\Control\HTTPResponse;
 use SilverStripe\Core\Config\Config;
@@ -12,6 +13,7 @@ use SilverStripe\Security\BasicAuth;
 use SilverStripe\Security\Member;
 use SilverStripe\Security\Security;
 use SilverStripe\Security\SecurityToken;
+use SilverStripe\Versioned\Versioned;
 use SilverStripe\View\SSViewer;
 use PHPUnit_Framework_AssertionFailedError;
 use SimpleXMLElement;
@@ -161,6 +163,9 @@ class FunctionalTest extends SapphireTest implements TestOnly
     public function get($url, $session = null, $headers = null, $cookies = null)
     {
         $this->cssParser = null;
+        if (Versioned::get_stage() === Versioned::DRAFT) {
+            $url = HTTP::setGetVar('stage', Versioned::DRAFT, $url);
+        }
         $response = $this->mainSession->get($url, $session, $headers, $cookies);
         if ($this->autoFollowRedirection && is_object($response) && $response->getHeader('Location')) {
             $response = $this->mainSession->followRedirection();
@@ -411,11 +416,9 @@ class FunctionalTest extends SapphireTest implements TestOnly
     public function useDraftSite($enabled = true)
     {
         if ($enabled) {
-            $this->session()->set('readingMode', 'Stage.Stage');
-            $this->session()->set('unsecuredDraftSite', true);
+            Versioned::set_stage(Versioned::DRAFT);
         } else {
-            $this->session()->set('readingMode', 'Stage.Live');
-            $this->session()->set('unsecuredDraftSite', false);
+            Versioned::set_stage(Versioned::LIVE);
         }
     }