diff --git a/README.md b/README.md
index db9283b..25e8f98 100644
--- a/README.md
+++ b/README.md
@@ -61,4 +61,3 @@ Parameters for "dev/testsession/start":
 Example usage with parameters:
 
 	dev/testsession/start?database=ss_tmpdb_1234567&fixture=cms/tests/controller/CMSMainTest.yml
-
diff --git a/code/TestSessionMemberExtension.php b/code/TestSessionMemberExtension.php
index e209bd7..ead8a6f 100644
--- a/code/TestSessionMemberExtension.php
+++ b/code/TestSessionMemberExtension.php
@@ -10,6 +10,22 @@ class TestSessionMemberExtension extends DataExtension {
 	public function memberLoggedIn() {
 		if(!SapphireTest::using_temp_db()) return;
 
+		$this->setCurrentMemberState();
+	}
+
+	public function onRegister() {
+		if(!SapphireTest::using_temp_db()) return;
+
+		$this->setCurrentMemberState();
+	}
+
+	public function memberLoggedOut() {
+		if(!SapphireTest::using_temp_db()) return;
+
+		$state = TestSessionDatabaseState::get()->filter('Key', 'CurrentMemberID')->removeAll();
+	}
+
+	protected function setCurrentMemberState() {
 		$state = TestSessionDatabaseState::get()->find('Key', 'CurrentMemberID');
 		if(!$state) {
 			$state = new TestSessionDatabaseState(array(
@@ -20,10 +36,4 @@ class TestSessionMemberExtension extends DataExtension {
 		$state->write();
 	}
 
-	public function memberLoggedOut() {
-		if(!SapphireTest::using_temp_db()) return;
-
-		$state = TestSessionDatabaseState::get()->filter('Key', 'CurrentMemberID')->removeAll();
-	}
-
 }
\ No newline at end of file