diff --git a/masquerade.install b/masquerade.install
index a10d092..8075f3b 100644
--- a/masquerade.install
+++ b/masquerade.install
@@ -2,15 +2,11 @@
 
 /**
  * @file
- * masquerade.install
- *
  * Install, uninstall and update hooks for the Masquarade module.
  */
 
 /**
  * Implements hook_schema().
- *
- * @return array
  */
 function masquerade_schema() {
   return array(
diff --git a/masquerade.module b/masquerade.module
index 56a0d78..5b279b4 100644
--- a/masquerade.module
+++ b/masquerade.module
@@ -19,8 +19,6 @@ function masquerade_help($path, $arg) {
 
 /**
  * Implements hook_permission().
- *
- * @return array
  */
 function masquerade_permission() {
   return array(
@@ -84,8 +82,6 @@ function masquerade_cron() {
  * Implements hook_menu().
  */
 function masquerade_menu() {
-  $items = array();
-
   $default_test_user = _masquerade_user_load(variable_get('masquerade_test_user', ''));
   if ($default_test_user && ($default_test_user->uid || $default_test_user->name == variable_get('anonymous', t('Anonymous')))) {
     $items['masquerade/switch/' . $default_test_user->uid] = array(
@@ -95,7 +91,6 @@ function masquerade_menu() {
       'page arguments' => array(2),
       'access callback' => 'masquerade_menu_access',
       'access arguments' => array('switch'),
-      'type' => MENU_NORMAL_ITEM,
     );
   }
 
@@ -105,14 +100,12 @@ function masquerade_menu() {
     'page arguments' => array(2),
     'access callback' => 'masquerade_menu_access',
     'access arguments' => array('switch', 2),
-    'type' => MENU_NORMAL_ITEM,
   );
   $items['masquerade/unswitch'] = array(
     'title' => 'Switch back',
     'page callback' => 'masquerade_switch_back_page',
     'access callback' => 'masquerade_menu_access',
     'access arguments' => array('unswitch'),
-    'type' => MENU_NORMAL_ITEM,
   );
   $items['masquerade/autocomplete'] = array(
     'title' => '',
@@ -142,7 +135,6 @@ function masquerade_menu() {
     'page arguments' => array('masquerade_admin_settings'),
     'access callback' => 'user_access',
     'access arguments' => array('administer masquerade'),
-    'type' => MENU_NORMAL_ITEM,
   );
 
   return $items;
@@ -222,13 +214,14 @@ function masquerade_menu_access($type, $uid = NULL) {
   switch ($type) {
     case 'unswitch':
       return isset($_SESSION['masquerading']) || arg(2) == 'menu-customize' || arg(2) == 'menu';
+
     case 'autocomplete':
       return isset($_SESSION['masquerading']) || (user_access('masquerade as user') || user_access('masquerade as admin'));
-      break;
+
     case 'user':
       global $user;
       return db_query("SELECT 1 FROM {masquerade_users} WHERE uid_from = :uid_from", array(':uid_from' => $user->uid))->fetchField();
-      break;
+
     case 'switch':
       $switch_to_account = FALSE;
       global $user;
@@ -244,7 +237,6 @@ function masquerade_menu_access($type, $uid = NULL) {
         }
       }
       return !isset($_SESSION['masquerading']) && (user_access('masquerade as user') || user_access('masquerade as admin') || $switch_to_account);
-      break;
   }
 }
 
@@ -252,8 +244,7 @@ function masquerade_menu_access($type, $uid = NULL) {
  * Admin settings form.
  */
 function masquerade_admin_settings() {
-  drupal_set_message(user_access('masquerade as user') ? 'Allowed' : 'Disallowed');
-  // create a list of roles; all selected roles are considered administrative.
+  // Create a list of roles; all selected roles are considered administrative.
   $roles = user_roles();
   $form['masquerade_admin_roles'] = array(
     '#type' => 'checkboxes',
@@ -360,8 +351,6 @@ function _masquerade_user_load($username) {
 function masquerade_user_logout($account) {
   if (!empty($account->masquerading)) {
     global $user;
-    // @TODO Seems this useless.
-    cache_clear_all($user->uid, 'cache_menu', TRUE);
     $real_user = user_load($user->masquerading);
     watchdog('masquerade', "User %user no longer masquerading as %masq_as.", array('%user' => $real_user->name, '%masq_as' => $user->name), WATCHDOG_INFO);
 
@@ -749,6 +738,10 @@ function masquerade_autocomplete_multiple($string, $add_anonymous = TRUE) {
  */
 function masquerade_switch_user_page($uid) {
   if (isset($_GET['token']) && drupal_valid_token($_GET['token'], 'masquerade/switch/' . $uid) && masquerade_switch_user($uid)) {
+    global $user;
+    drupal_set_message(t('You are now masquerading as !masq_as.', array(
+      '!masq_as' => theme('username', array('account' => $user)),
+    )));
     drupal_goto($_SERVER['HTTP_REFERER']);
   }
   else {
@@ -759,11 +752,11 @@ function masquerade_switch_user_page($uid) {
 /**
  * Allows a user with the right permissions to become the selected user.
  *
- * @param $uid
+ * @param int $uid
  *   The user ID to switch to.
  *
- * @return
- *   TRUE if the user was sucessfully switched, or FALSE if there was an error.
+ * @return bool
+ *   TRUE if the user was successfully switched, or FALSE if there was an error.
  */
 function masquerade_switch_user($uid) {
   global $user;
@@ -813,8 +806,10 @@ function masquerade_switch_user($uid) {
     ))
     ->execute();
 
-  watchdog('masquerade', 'User %user now masquerading as %masq_as.', array('%user' => $user->name, '%masq_as' => $new_user->name ? $new_user->name : variable_get('anonymous', t('Anonymous'))), WATCHDOG_INFO);
-  drupal_set_message(t('You are now masquerading as !masq_as.', array('!masq_as' => theme('username', array('account' => $new_user)))));
+  watchdog('masquerade', 'User %user now masquerading as %masq_as.', array(
+    '%user' => $user->name,
+    '%masq_as' => $new_user->name ? $new_user->name : variable_get('anonymous', t('Anonymous')),
+  ), WATCHDOG_INFO);
   $new_user->masquerading = $user->uid;
   $user = $new_user;
 
@@ -831,9 +826,12 @@ function masquerade_switch_user($uid) {
 function masquerade_switch_back_page() {
   if (isset($_GET['token']) && drupal_valid_token($_GET['token'], 'masquerade/unswitch')) {
     global $user;
-    $olduser = $user;
+    $old_user = $user;
     masquerade_switch_back();
-    drupal_set_message(t('You are no longer masquerading as !masq_as and are now logged in as !user.', array('!user' => theme('username', array('account' => $user)), '!masq_as' => theme('username', array('account' => $olduser)))));
+    drupal_set_message(t('You are no longer masquerading as !masq_as and are now logged in as !user.', array(
+      '!user' => theme('username', array('account' => $user)),
+      '!masq_as' => theme('username', array('account' => $old_user)),
+    )));
     drupal_goto($_SERVER['HTTP_REFERER']);
   }
   else {
@@ -847,8 +845,6 @@ function masquerade_switch_back_page() {
 function masquerade_switch_back() {
   // Switch user.
   global $user;
-  // @TODO Seems useless.
-  cache_clear_all($user->uid, 'cache_menu', TRUE);
   $uid = db_query("SELECT m.uid_from FROM {masquerade} m WHERE m.sid = :sid AND m.uid_as = :uid_as ", array(
     ':sid' => session_id(),
     ':uid_as' => $user->uid,
diff --git a/masquerade.test b/masquerade.test
index c1a6fa8..84da8ff 100644
--- a/masquerade.test
+++ b/masquerade.test
@@ -2,8 +2,6 @@
 
 /**
  * @file
- * masquerade.test
- *
  * Test the form permissions and switch ability of the Masquarade module.
  */
 
@@ -43,15 +41,34 @@ class MasqueradeTestCase extends DrupalWebTestCase {
 
     // Tests switch.
     $switch_link = 'masquerade/switch/' . $web_user->uid;
-    debug($this->drupalGetToken($switch_link), 'drupalGetToken');
-    debug( drupal_hmac_base64($switch_link, $this->session_id . drupal_get_private_key(). drupal_get_hash_salt()));
-    $this->drupalGet($switch_link, array('query' => array('token' => $this->drupalGetToken($switch_link))));
-    $this->assertText('Now masquerading as ' . $web_user->name);
+    // Tests token calculation.
+    $switch_token = drupal_hmac_base64($switch_link, $this->session_id . drupal_get_private_key(). drupal_get_hash_salt());
+    // @todo Uncomment after http://drupal.org/node/1555862 lands.
+    //$this->assertEqual($this->drupalGetToken($switch_link), $switch_token, 'Calculated token equals drupalGetToken()');
+    //$this->drupalGet($switch_link, array('query' => array('token' => $this->drupalGetToken($switch_link))));
+    // @todo Remove after http://drupal.org/node/1555862 lands.
+    $this->drupalGet($switch_link, array('query' => array('token' => $switch_token)));
+
+    //$this->assertText('Now masquerading as ' . $web_user->name);
+    $this->assertText(format_string('You are now masquerading as @masq_as.', array(
+      '@masq_as' => $web_user->name,
+    )));
 
     // Tests switch back.
     $unswitch_link = 'masquerade/unswitch';
-    $this->drupalGet($unswitch_link, array('query' => array('token' => $this->drupalGetToken($unswitch_link))));
-    $this->assertText('No longer masquerading as ' . $web_user->name);
+    $unswitch_token = drupal_hmac_base64($unswitch_link, $this->session_id . drupal_get_private_key(). drupal_get_hash_salt());
+    // @todo Uncomment after http://drupal.org/node/1555862 lands.
+    //$this->assertEqual($this->drupalGetToken($unswitch_link), $unswitch_token, 'Calculated token equals drupalGetToken()');
+    //$this->drupalGet($unswitch_link, array('query' => array('token' => $this->drupalGetToken($unswitch_link))));
+    // @todo Remove after http://drupal.org/node/1555862 lands.
+    $this->drupalGet($unswitch_link, array('query' => array('token' => $unswitch_token)));
+
+    //$this->assertText('No longer masquerading as ' . $web_user->name);
+    $this->assertText(format_string('You are no longer masquerading as @masq_as and are now logged in as @user.', array(
+      '@masq_as' => $web_user->name,
+      '@user' => $admin_user->name,
+    )));
+
   }
 
 }