diff --git a/src/Dev/FunctionalTest.php b/src/Dev/FunctionalTest.php
index 2f3d6a6dc..45f0ce020 100644
--- a/src/Dev/FunctionalTest.php
+++ b/src/Dev/FunctionalTest.php
@@ -403,16 +403,17 @@ class FunctionalTest extends SapphireTest implements TestOnly
      * This is helpful if you're not testing publication functionality and don't want "stage management" cluttering
      * your test.
      *
-     * @param bool $draft toggle the use of the draft site
+     * @param bool $enabled toggle the use of the draft site
      */
-    public function useDraftSite($draft = true)
+    public function useDraftSite($enabled = true)
     {
-        if (!class_exists(Versioned::class)) {
-            return;
+        if ($enabled) {
+            $this->session()->set('readingMode', 'Stage.Stage');
+            $this->session()->set('unsecuredDraftSite', true);
+        } else {
+            $this->session()->set('readingMode', 'Stage.Live');
+            $this->session()->set('unsecuredDraftSite', false);
         }
-        $stage = $draft ? Versioned::DRAFT : Versioned::LIVE;
-        Versioned::set_stage($stage);
-        Versioned::set_default_reading_mode(Versioned::get_reading_mode());
     }
 
     /**
diff --git a/src/Dev/TestSession.php b/src/Dev/TestSession.php
index 171ec509f..a27142632 100644
--- a/src/Dev/TestSession.php
+++ b/src/Dev/TestSession.php
@@ -2,15 +2,16 @@
 
 namespace SilverStripe\Dev;
 
-use SilverStripe\Control\Cookie_Backend;
-use SilverStripe\Control\HTTPRequest;
-use SilverStripe\Control\Session;
+use Exception;
 use SilverStripe\Control\Controller;
+use SilverStripe\Control\Cookie_Backend;
 use SilverStripe\Control\Director;
+use SilverStripe\Control\HTTPRequest;
 use SilverStripe\Control\HTTPResponse;
+use SilverStripe\Control\Session;
+use SilverStripe\Core\Extensible;
 use SilverStripe\Core\Injector\Injector;
 use SimpleByName;
-use Exception;
 use SimplePage;
 use SimplePageBuilder;
 
@@ -20,6 +21,7 @@ use SimplePageBuilder;
  */
 class TestSession
 {
+    use Extensible;
 
     /**
      * @var Session
@@ -89,6 +91,7 @@ class TestSession
      */
     public function get($url, $session = null, $headers = null, $cookies = null)
     {
+        $this->extend('updateGetURL', $url, $session, $headers, $cookies);
         $headers = (array) $headers;
         if ($this->lastUrl && !isset($headers['Referer'])) {
             $headers['Referer'] = $this->lastUrl;
@@ -123,6 +126,7 @@ class TestSession
      */
     public function post($url, $data, $headers = null, $session = null, $body = null, $cookies = null)
     {
+        $this->extend('updatePostURL', $url, $data, $headers, $session, $body, $cookies);
         $headers = (array) $headers;
         if ($this->lastUrl && !isset($headers['Referer'])) {
             $headers['Referer'] = $this->lastUrl;