diff --git a/composer.json b/composer.json
index 765a14a..dccb639 100644
--- a/composer.json
+++ b/composer.json
@@ -38,7 +38,6 @@
 			"SilverStripe\\BehatExtension\\Tests": "tests/"
 		},
 		"classmap": [
-			"testsession/code",
 			"framework",
 			"vendor/phpunit/phpunit"
 		]
diff --git a/src/SilverStripe/BehatExtension/Context/Initializer/SilverStripeAwareInitializer.php b/src/SilverStripe/BehatExtension/Context/Initializer/SilverStripeAwareInitializer.php
index 2c7f56e..5afbfe1 100644
--- a/src/SilverStripe/BehatExtension/Context/Initializer/SilverStripeAwareInitializer.php
+++ b/src/SilverStripe/BehatExtension/Context/Initializer/SilverStripeAwareInitializer.php
@@ -7,6 +7,7 @@ use Behat\Behat\Context\ContextInterface;
 use SilverStripe\BehatExtension\Context\SilverStripeAwareContextInterface;
 use SilverStripe\Core\Injector\Injector;
 use SilverStripe\Dev\SapphireTest;
+use SilverStripe\TestSession\TestSessionEnvironment;
 
 /*
  * This file is part of the Behat/SilverStripeExtension
@@ -74,7 +75,7 @@ class SilverStripeAwareInitializer implements InitializerInterface
 
         file_put_contents('php://stdout', "Creating test session environment" . PHP_EOL);
 
-        $testEnv = Injector::inst()->get('TestSessionEnvironment');
+        $testEnv = TestSessionEnvironment::singleton();
         $testEnv->startTestSession(array(
             'createDatabase' => true
         ));
diff --git a/src/SilverStripe/BehatExtension/Context/SilverStripeContext.php b/src/SilverStripe/BehatExtension/Context/SilverStripeContext.php
index 4f966f9..27f5c8e 100644
--- a/src/SilverStripe/BehatExtension/Context/SilverStripeContext.php
+++ b/src/SilverStripe/BehatExtension/Context/SilverStripeContext.php
@@ -12,6 +12,8 @@ use Behat\Mink\Exception\ElementNotFoundException;
 use InvalidArgumentException;
 use SilverStripe\BehatExtension\Context\SilverStripeAwareContextInterface;
 use Symfony\Component\Yaml\Yaml;
+use SilverStripe\TestSession\TestSessionEnvironment;
+
 
 // Mink etc.
 require_once 'vendor/autoload.php';
@@ -72,7 +74,7 @@ class SilverStripeContext extends MinkContext implements SilverStripeAwareContex
     {
         // Initialize your context here
         $this->context = $parameters;
-        $this->testSessionEnvironment = new \TestSessionEnvironment();
+        $this->testSessionEnvironment = new TestSessionEnvironment();
     }
 
     public function setDatabase($databaseName)
diff --git a/src/SilverStripe/BehatExtension/Utility/TestMailer.php b/src/SilverStripe/BehatExtension/Utility/TestMailer.php
index dfbb2f0..1e1cbe7 100644
--- a/src/SilverStripe/BehatExtension/Utility/TestMailer.php
+++ b/src/SilverStripe/BehatExtension/Utility/TestMailer.php
@@ -4,7 +4,7 @@ namespace SilverStripe\BehatExtension\Utility;
 
 use SilverStripe\Dev\TestMailer as BaseTestMailer;
 use SilverStripe\Core\Injector\Injector;
-use TestSessionEnvironment;
+use SilverStripe\TestSession\TestSessionEnvironment;
 
 /**
  * Same principle as core TestMailer class,
@@ -21,7 +21,7 @@ class TestMailer extends BaseTestMailer
 
     public function __construct()
     {
-        $this->testSessionEnvironment = Injector::inst()->get('TestSessionEnvironment');
+        $this->testSessionEnvironment = TestSessionEnvironment::singleton();
     }
 
     /**