From e957d1e0fd5697d8fe30e13d968f1e394efb0ec3 Mon Sep 17 00:00:00 2001
From: UndefinedOffset <echipman@webbuildersgroup.com>
Date: Thu, 24 Jan 2019 15:22:29 -0400
Subject: [PATCH] BUGFIX: Fixed issue where the incorrect database connection
 could be made when using a stubfile (fixes #60)

---
 src/TestSessionEnvironment.php    | 54 +++++++++++++++++++------------
 src/TestSessionHTTPMiddleware.php |  6 ++--
 2 files changed, 36 insertions(+), 24 deletions(-)

diff --git a/src/TestSessionEnvironment.php b/src/TestSessionEnvironment.php
index f6fcb09..fe42208 100644
--- a/src/TestSessionEnvironment.php
+++ b/src/TestSessionEnvironment.php
@@ -284,26 +284,7 @@ class TestSessionEnvironment
         }
 
         // ensure we have a connection to the database
-        if (isset($state->database) && $state->database) {
-            if (!DB::get_conn()) {
-                // No connection, so try and connect to tmpdb if it exists
-                if (isset($state->database)) {
-                    $this->oldDatabaseName = $databaseConfig['database'];
-                    $databaseConfig['database'] = $state->database;
-                }
-
-                // Connect to database
-                DB::connect($databaseConfig);
-            } else {
-                // We've already connected to the database, do a fast check to see what database we're currently using
-                $db = DB::get_conn()->getSelectedDatabase();
-                if (isset($state->database) && $db != $state->database) {
-                    $this->oldDatabaseName = $databaseConfig['database'];
-                    $databaseConfig['database'] = $state->database;
-                    DB::connect($databaseConfig);
-                }
-            }
-        }
+        $this->connectToDatabase($state);
 
         // Database
         if (!$this->isRunningTests()) {
@@ -562,6 +543,39 @@ class TestSessionEnvironment
         return PUBLIC_PATH . DIRECTORY_SEPARATOR . 'assets_backup';
     }
 
+    /**
+     * Ensure that there is a connection to the database
+     * 
+     * @param mixed $state
+     */
+    public function connectToDatabase($state = null) {
+        if ($state == null) {
+            $state = $this->getState();
+        }
+
+        $databaseConfig = DB::getConfig();
+
+        if (isset($state->database) && $state->database) {
+            if (!DB::get_conn()) {
+                // No connection, so try and connect to tmpdb if it exists
+                if (isset($state->database)) {
+                    $this->oldDatabaseName = $databaseConfig['database'];
+                    $databaseConfig['database'] = $state->database;
+                }
+
+                // Connect to database
+                DB::connect($databaseConfig);
+            } else {
+                // We've already connected to the database, do a fast check to see what database we're currently using
+                $db = DB::get_conn()->getSelectedDatabase();
+                if (isset($state->database) && $db != $state->database) {
+                    $this->oldDatabaseName = $databaseConfig['database'];
+                    $databaseConfig['database'] = $state->database;
+                    DB::connect($databaseConfig);
+                }
+            }
+        }
+    }
 
     /**
      * Wait for pending requests
diff --git a/src/TestSessionHTTPMiddleware.php b/src/TestSessionHTTPMiddleware.php
index ac8be5c..6dd69e0 100644
--- a/src/TestSessionHTTPMiddleware.php
+++ b/src/TestSessionHTTPMiddleware.php
@@ -80,10 +80,8 @@ class TestSessionHTTPMiddleware implements HTTPMiddleware
             $file = $testState->stubfile;
             if (!Director::isLive() && $file && file_exists($file)) {
                 // Connect to the database so the included code can interact with it
-                $databaseConfig = DB::getConfig();
-                if ($databaseConfig) {
-                    DB::connect($databaseConfig);
-                }
+                $this->testSessionEnvironment->connectToDatabase();
+
                 include_once($file);
             }
         }