diff --git a/README.txt b/README.txt
index a406cae..df420b8 100644
--- a/README.txt
+++ b/README.txt
@@ -31,7 +31,7 @@ CONTENTS
 ----
 1.  Introduction
 
-Workbench Moderation 
+Workbench Moderation
 
 ----
 1.1  Concepts
@@ -129,6 +129,13 @@ them where your users can find them.
 Using the "Workbench" module with Workbench Moderation enables the display of
 moderation status information and a mini moderation form on node viewing pages.
 
+There is one dependency:
+
+  https://www.drupal.org/project/drafty
+
+The Drafty module is used for managing changes to the node's state and must also
+be installed.
+
 ----
 3.  Configuration
 
@@ -176,7 +183,7 @@ perform a particular moderation task.
 ----
 3.3.1  Recommended permissions
 
-For reference, these are the permission sets recommended by the "Check 
+For reference, these are the permission sets recommended by the "Check
 Permissions" tab:
 
     Author:
@@ -189,7 +196,7 @@ Permissions" tab:
       Workbench Moderation:
         view moderation messages
         use workbench_moderation my drafts tab
-    
+
     Editor:
       Node:
         access content
@@ -203,7 +210,7 @@ Permissions" tab:
         view moderation history
         use workbench_moderation my drafts tab
         use workbench_moderation needs review tab
-    
+
     Moderator:
       Node:
         access content
diff --git a/tests/external_node_update.test b/tests/external_node_update.test
index 9edbcf8..fd41ee0 100644
--- a/tests/external_node_update.test
+++ b/tests/external_node_update.test
@@ -178,12 +178,13 @@ class WorkbenchModerationExternalNodeUpdateTestCase extends WorkbenchModerationT
    *   The node's record(s) from the {workbench_moderation_node_history} table.
    */
   protected function getModerationRecord($status = 'is_current') {
-    return db_select('workbench_moderation_node_history', 'nh')
+    $data = db_select('workbench_moderation_node_history', 'nh')
       ->fields('nh', array('from_state', 'state', 'published', 'is_current'))
       ->condition('nid', $this->node->nid, '=')
       ->condition($status, 1)
       ->execute()
       ->fetchAssoc();
+    return $data;
   }
 
 }
diff --git a/tests/workbench_moderation.node_access.test b/tests/workbench_moderation.node_access.test
new file mode 100644
index 0000000..92dd196
--- /dev/null
+++ b/tests/workbench_moderation.node_access.test
@@ -0,0 +1,133 @@
+<?php
+
+/**
+ * @file
+ * Node Access tests for workbench_moderation.module.
+ */
+
+class WorkbenchModerationNodeAccessTestCase extends WorkbenchModerationTestCase {
+  protected $test_user;
+
+  public static function getInfo() {
+    return array(
+      'name' => 'Node Access test',
+      'description' => 'Ensure that node access rules are enforced during moderation.',
+      'group' => 'Workbench Moderation',
+    );
+  }
+
+  function setUp($modules = array()) {
+    $modules = array_merge($modules, array('workbench_moderation', 'workbench_moderation_node_access_test'));
+    parent::setUp($modules);
+  }
+
+  /**
+   * Creates a node and tests the creation of node access rules.
+   */
+  function testNodeAccessRecords() {
+    // Make sure this node type is moderated.
+    $is_moderated = workbench_moderation_node_type_moderated($this->content_type);
+    $this->assertTrue($is_moderated, t('The content type is moderated.'));
+
+    // Create an published node that has access records.
+    $node1 = $this->drupalCreateNode(array('type' => $this->content_type, 'title' => 'Published node', 'status' => 1));
+    $this->assertTrue(node_load($node1->nid), 'Test published node created.');
+
+    // Assert grants were added to node_access on creation.
+    $records = db_query('SELECT realm, gid FROM {node_access} WHERE nid = :nid', array(':nid' => $node1->nid))->fetchAll();
+    $this->assertEqual(count($records), 1, 'Returned the correct number of rows.');
+    $this->assertEqual($records[0]->realm, 'test_wm_realm', 'Grant with test_wm_realm acquired for node without alteration.');
+    $this->assertEqual($records[0]->gid, 1, 'Grant with gid = 1 acquired for node without alteration.');
+
+    // Create an unpublished node that has no access records.
+    $node2 = $this->drupalCreateNode(array('type' => $this->content_type, 'title' => 'Unpublished node', 'status' => 0));
+    $this->assertTrue(node_load($node2->nid), 'Test unpublished node created.');
+
+    // Check that no access records are present for the node, since it is not published.
+    $records = db_query('SELECT realm, gid FROM {node_access} WHERE nid = :nid', array(':nid' => $node2->nid))->fetchAll();
+    $this->assertEqual(count($records), 0, 'Returned no records for unpublished node.');
+
+    // Login for form testing.
+    $this->drupalLogin($this->moderator_user);
+
+    // Create a new draft of the published node.
+    $edit = array('title' => 'Draft node revision1');
+    $this->drupalPost("node/{$node1->nid}/edit", $edit, t('Save'));
+
+    // Assert grants still exist on node_access after draft created.
+    $records = db_query('SELECT realm, gid FROM {node_access} WHERE nid = :nid', array(':nid' => $node1->nid))->fetchAll();
+    $this->assertEqual(count($records), 1, 'Returned the correct number of rows.');
+    $this->assertEqual($records[0]->realm, 'test_wm_realm', 'Grant with article_realm acquired for node without alteration.');
+    $this->assertEqual($records[0]->gid, 1, 'Grant with gid = 1 acquired for node without alteration.');
+
+    // Create an published node that has access records.
+    $node3 = $this->drupalCreateNode(array('type' => $this->content_type, 'title' => 'Published node', 'status' => 1));
+    $this->assertTrue(node_load($node3->nid), 'Test published node created.');
+
+    // Make sure the "Moderate" tab is accessible.
+    $this->drupalGet("node/{$node3->nid}/moderation");
+
+    // Attempt to change the moderation state without a token in the link.
+    $this->drupalGet("node/{$node3->nid}/moderation/{$node3->vid}/change-state/needs_review");
+    $this->assertResponse(403);
+
+    // Run the same state change with a valid token, which should succeed.
+    $this->drupalGet("node/{$node3->nid}/moderation/{$node3->vid}/change-state/needs_review", array(
+      'query' => array('token' => $this->drupalGetToken("{$node3->nid}:{$node3->vid}:needs_review"))
+    ));
+
+    // Test that we have unpublished the node.
+    $this->assertResponse(200);
+    $node = node_load($node3->nid, NULL, TRUE);
+    $this->assertEqual($node->workbench_moderation['current']->state, 'needs_review', 'Node state changed to Needs Review via callback.');
+    $this->assertFalse($node->status, 'Node is no longer published after moderation.');
+
+    // Check that no access records are present for the node, since it is not published.
+    $records = db_query('SELECT realm, gid FROM {node_access} WHERE nid = :nid', array(':nid' => $node3->nid))->fetchAll();
+    $this->assertEqual(count($records), 0, 'Returned no records for unpublished node.');
+
+    // Publish the node via the moderation form.
+    $moderate = array('state' => workbench_moderation_state_published());
+    $this->drupalPost("node/{$node3->nid}/moderation", $moderate, t('Apply'));
+
+    // Ensure published status.
+    $node = node_load($node3->nid, NULL, TRUE);
+    $this->assertTrue(isset($node->workbench_moderation['published']), t('Workbench moderation has a published revision'));
+
+    // Check to see if grants are still present.
+    $records = db_query('SELECT realm, gid FROM {node_access} WHERE nid = :nid', array(':nid' => $node3->nid))->fetchAll();
+    $this->assertEqual(count($records), 1, 'Returned the correct number of rows.');
+    $this->assertEqual($records[0]->realm, 'test_wm_realm', 'Grant with article_realm acquired for node without alteration.');
+    $this->assertEqual($records[0]->gid, 1, 'Grant with gid = 1 acquired for node without alteration.');
+
+    // Create a new draft.
+    $edit = array('title' => $this->randomName(10) . '_revision1');
+    $this->drupalPost("node/{$node3->nid}/edit", $edit, t('Save'));
+
+    // Check that grants are not changed.
+    $records = db_query('SELECT realm, gid FROM {node_access} WHERE nid = :nid', array(':nid' => $node3->nid))->fetchAll();
+    $this->assertEqual(count($records), 1, 'Returned the correct number of rows.');
+    $this->assertEqual($records[0]->realm, 'test_wm_realm', 'Grant with article_realm acquired for node without alteration.');
+    $this->assertEqual($records[0]->gid, 1, 'Grant with gid = 1 acquired for node without alteration.');
+
+    // Ensure that the published version is still present.
+    $node = node_load($node3->nid, NULL, TRUE);
+    debug($node->title);
+    $this->assertTrue($node->status, t('Content is published'));
+    $this->assertEqual($node->title, 'Published node', t('Published version is loaded.'));
+
+    // Load the published and draft revisions.
+    $draft = clone $node;
+    $draft = workbench_moderation_node_current_load($draft);
+
+    // Assert the two revisions are unique.
+    $this->assertEqual($node->vid, $node->workbench_moderation['published']->vid, t('Published revision is loaded by default'));
+    $this->assertTrue($node->status, t('Published revision has status = 1'));
+    $this->assertNotEqual($node->vid, $draft->vid, t('Draft revision is different from the published revision'));
+
+    // Test the node for an anon user.
+    $this->drupalLogout();
+    $this->drupalGet('node/' . $node->nid);
+    $this->assertRaw('Published node', t('Page title is "Published node"'));
+  }
+}
diff --git a/tests/workbench_moderation.perms.test b/tests/workbench_moderation.perms.test
new file mode 100644
index 0000000..e82242b
--- /dev/null
+++ b/tests/workbench_moderation.perms.test
@@ -0,0 +1,82 @@
+<?php
+
+/**
+ * @file
+ * Permissions-related tests for workbench_moderation.module.
+ */
+
+class WorkbenchModerationPermsTestCase extends DrupalWebTestCase {
+  protected $content_type;
+  protected $editor_user;
+  protected $author_user;
+
+  function setUp($modules = array()) {
+    $modules = array_merge($modules, array('workbench_moderation'));
+    parent::setUp($modules);
+
+    // Create a new content type and enable moderation on it.
+    $type = $this->drupalCreateContentType();
+    $this->content_type = $type->name;
+    variable_set('node_options_' . $this->content_type, array('revision', 'moderation'));
+    // The editor should be able to view all unpublished content, even without authoring perms.
+    $editor_permissions = array(
+      0 => 'view all unpublished content',
+    );
+    $this->editor_user = $this->drupalCreateUser($editor_permissions);
+    // The Author will create the content.
+    $author_permissions = array(
+      0 => 'create ' . $type->name . ' content',
+    );
+    $this->author_user = $this->drupalCreateUser($author_permissions);
+  }
+}
+
+class WorkbenchModerationViewUnpublishedTestCase extends WorkbenchModerationPermsTestCase {
+
+  public static function getInfo() {
+    return array(
+      'name' => 'View all unpublished content',
+      'description' => 'Create a user who can view unpublished content. Create a node and leave it unpublished. Try to view it.',
+      'group' => 'Workbench Moderation',
+    );
+  }
+
+  function setUp($modules = array()) {
+    parent::setUp($modules);
+    $this->drupalLogin($this->author_user);
+  }
+
+  function testViewUnpublished() {
+    $is_moderated = workbench_moderation_node_type_moderated($this->content_type);
+    $this->assertTrue($is_moderated, t('The content type is moderated.'));
+
+    // Create a new node and make sure it is unpublished.
+    $body_name = 'body[' . LANGUAGE_NONE . '][0]';
+    $edit = array(
+      'title' => $this->randomName(),
+      "{$body_name}[value]" => $this->randomString(128),
+      "{$body_name}[format]" => filter_default_format(),
+    );
+    $this->drupalPost("node/add/{$this->content_type}", $edit, t('Save'));
+
+    // Get the new node.
+    $node = $this->drupalGetNodeByTitle($edit['title']);
+
+    $this->assertFalse($node->status, t('New node is unpublished'));
+    $this->assertTrue(isset($node->workbench_moderation), t('Workbench moderation information is present on the node object'));
+    $this->assertFalse(isset($node->workbench_moderation['published']), t('Workbench moderation has no published revision'));
+    $this->assertEqual($node->uid, $this->author_user->uid, 'This node was authored by the author user.');
+    $this->verbose(print_r($this->loggedInUser, TRUE));
+
+    $this->drupalLogin($this->editor_user);
+    global $user;
+    $user = user_load($this->loggedInUser->uid);
+    $this->drupalGet($node->path['source']);
+    $this->assertFalse($node->status, t('This node is unpublished.'));
+    $this->assertResponse(200);
+    $this->assertFalse($node->uid == $this->loggedInUser->uid, t('The current user is not the author of this node.'));
+    $this->assertEqual($user->uid, $this->loggedInUser->uid, 'The current global user is the same as the logged in user.');
+    $this->assertEqual($user->uid, $this->editor_user->uid, 'The current user is the editor user.');
+    $this->assertTrue(user_access('view all unpublished content'), 'Current user has permission to view all unpublished content');
+  }
+}
diff --git a/tests/workbench_moderation.test b/tests/workbench_moderation.test
index 42d1665..4559a57 100644
--- a/tests/workbench_moderation.test
+++ b/tests/workbench_moderation.test
@@ -41,6 +41,17 @@ class WorkbenchModerationTestCase extends DrupalWebTestCase {
     ));
   }
 
+  /**
+   * Override DrupalWebTestCase::drupalGetToken() as it does not return the
+   * correct token for the currently logged-in testing user.
+   */
+  protected function drupalGetToken($value = '') {
+    $session_id = $this->session_id;
+    if (empty($session_id) && !empty($this->loggedInUser)) {
+      $session_id = db_query("SELECT sid FROM {sessions} WHERE uid = :uid ORDER BY timestamp DESC", array(':uid' => $this->loggedInUser->uid))->fetchField();
+    }
+    return drupal_hmac_base64($value, $session_id . drupal_get_private_key() . drupal_get_hash_salt());
+  }
 }
 
 class WorkbenchModerationModerateTabTestCase extends WorkbenchModerationTestCase {
@@ -58,18 +69,6 @@ class WorkbenchModerationModerateTabTestCase extends WorkbenchModerationTestCase
     $this->drupalLogin($this->moderator_user);
   }
 
-  /**
-   * Override DrupalWebTestCase::drupalGetToken() as it does not return the
-   * correct token for the currently logged-in testing user.
-   */
-  protected function drupalGetToken($value = '') {
-    $session_id = $this->session_id;
-    if (empty($session_id) && !empty($this->loggedInUser)) {
-      $session_id = db_query("SELECT sid FROM {sessions} WHERE uid = :uid ORDER BY timestamp DESC", array(':uid' => $this->loggedInUser->uid))->fetchField();
-    }
-    return drupal_hmac_base64($value, $session_id . drupal_get_private_key() . drupal_get_hash_salt());
-  }
-
   function testModerateTab() {
     $is_moderated = workbench_moderation_node_type_moderated($this->content_type);
     $this->assertTrue($is_moderated, t('The content type is moderated.'));
diff --git a/tests/workbench_moderation.transition.test b/tests/workbench_moderation.transition.test
new file mode 100644
index 0000000..98ef8f1
--- /dev/null
+++ b/tests/workbench_moderation.transition.test
@@ -0,0 +1,224 @@
+<?php
+
+/**
+ * @file
+ * Tests for node transition hooks with workbench_moderation.module.
+ */
+
+class WorkbenchModerationTransitionTestCase extends WorkbenchModerationTestCase {
+  protected $properties;
+
+  public static function getInfo() {
+    return array(
+      'name' => 'Transition hook test',
+      'description' => 'Tests the usage of hook_workbench_moderation_transition().',
+      'group' => 'Workbench Moderation',
+    );
+  }
+
+  function setUp($modules = array()) {
+    // Enable a test module that will publish and unpublish nodes for us and
+    // provide hook implementations.
+    parent::setUp(array_merge($modules, array('workbench_moderation_test')));
+    $this->drupalLogin($this->moderator_user);
+
+    $this->properties = array('title', 'nid', 'vid', 'previous_state', 'new_state');
+  }
+
+  /**
+   * Asserts that transition values are as expected.
+   *
+   * @param array $expected
+   *   An array of values to be set by the transition.
+   */
+  private function assertTransition($expected = array()) {
+    foreach ($this->properties as $name) {
+      $value = workbench_moderation_test_get($name);
+      $this->assertEqual($value, $expected[$name], "Transition success: $name set to $value");
+    }
+  }
+
+  function testTransitionFromNodeForm() {
+    // Create a new node and publish it immediately.
+    $body_name = 'body[' . LANGUAGE_NONE . '][0]';
+    $edit = array(
+      'title' => $this->randomName(),
+      "{$body_name}[value]" => $this->randomString(128),
+      "{$body_name}[format]" => filter_default_format(),
+      'workbench_moderation_state_new' => workbench_moderation_state_published(),
+    );
+    $this->drupalPost("node/add/{$this->content_type}", $edit, t('Save'));
+
+    // Test the published state transition.
+    $expected = array(
+      'nid' => 1,
+      'vid' => 1,
+      'title' => $edit['title'],
+      'previous_state' => 'draft',
+      'new_state' => 'published',
+      'status' => 1,
+    );
+    $this->assertTransition($expected);
+
+    // Create a new draft of the published node.
+    $node = $this->drupalGetNodeByTitle($edit['title']);
+    $edit = array('title' => 'Draft node revision1');
+    $this->drupalPost("node/{$node->nid}/edit", $edit, t('Save'));
+
+    $expected = array(
+      'nid' => 1,
+      'vid' => 2,
+      'title' => $edit['title'],
+      'previous_state' => 'published',
+      'new_state' => 'draft',
+      'status' => 0,
+    );
+    $this->assertTransition($expected);
+
+    // Moderate to needs review and check transition.
+    $this->drupalGet("node/1/moderation/2/change-state/needs_review", array(
+      'query' => array('token' => $this->drupalGetToken("1:2:needs_review"))
+    ));
+    $expected = array(
+      'nid' => 1,
+      'vid' => 2,
+      'title' => $edit['title'],
+      'previous_state' => 'draft',
+      'new_state' => 'needs_review',
+      'status' => 0,
+    );
+    $this->assertTransition($expected);
+
+    // Publish the revision and check transition.
+    $this->drupalGet("node/1/moderation/2/change-state/published", array(
+      'query' => array('token' => $this->drupalGetToken("1:2:published"))
+    ));
+    $expected = array(
+      'nid' => 1,
+      'vid' => 2,
+      'title' => $edit['title'],
+      'previous_state' => 'needs_review',
+      'new_state' => 'published',
+      'status' => 1,
+    );
+    $this->assertTransition($expected);
+
+    // Create a new node and make sure it is unpublished.
+    $body_name = 'body[' . LANGUAGE_NONE . '][0]';
+    $edit = array(
+      'title' => $this->randomName(),
+      "{$body_name}[value]" => $this->randomString(128),
+      "{$body_name}[format]" => filter_default_format(),
+    );
+    $this->drupalPost("node/add/{$this->content_type}", $edit, t('Save'));
+
+    // Get the new node.
+    $node = $this->drupalGetNodeByTitle($edit['title']);
+
+    $expected = array(
+      'nid' => $node->nid,
+      'vid' => $node->vid,
+      'title' => $edit['title'],
+      'previous_state' => 'draft',
+      'new_state' => 'draft',
+      'status' => 0,
+    );
+    $this->assertTransition($expected);
+  }
+
+  function testUnpublishTransition() {
+    // Create a new node and publish it immediately. Assumes that
+    // WorkbenchModerationPublishFromNodeFormTestCase passes.
+    $body_name = 'body[' . LANGUAGE_NONE . '][0]';
+    $edit = array(
+      'title' => $this->randomName(),
+      "{$body_name}[value]" => $this->randomString(128),
+      "{$body_name}[format]" => filter_default_format(),
+      'workbench_moderation_state_new' => workbench_moderation_state_published(),
+    );
+    $this->drupalPost("node/add/{$this->content_type}", $edit, t('Save'));
+    $node = $this->drupalGetNodeByTitle($edit['title']);
+
+    $expected = array(
+      'nid' => 1,
+      'vid' => 1,
+      'title' => $edit['title'],
+      'previous_state' => 'draft',
+      'new_state' => 'published',
+      'status' => 1,
+    );
+    $this->assertTransition($expected);
+
+    // Unpublish the node via the unpublish confirmation form.
+    $this->drupalPost("node/{$node->nid}/moderation/{$node->vid}/unpublish", array(), t('Unpublish'));
+
+    $expected = array(
+      'nid' => 1,
+      'vid' => 1,
+      'title' => $edit['title'],
+      'previous_state' => 'published',
+      'new_state' => 'draft',
+      'status' => 0,
+    );
+    $this->assertTransition($expected);
+  }
+
+  public function testTransitionFromNodeSave() {
+    // Create a brand new unpublished node programmatically.
+    $edit = array(
+      'title' => $this->randomName(),
+      'type' => $this->content_type,
+      'status' => NODE_NOT_PUBLISHED,
+    );
+    $this->node = $this->drupalCreateNode($edit);
+
+    // Get the new node.
+    $node = $this->drupalGetNodeByTitle($edit['title']);
+
+    $expected = array(
+      'nid' => 1,
+      'vid' => 1,
+      'title' => $edit['title'],
+      'previous_state' => 'draft',
+      'new_state' => 'draft',
+      'status' => 0,
+    );
+    $this->assertTransition($expected);
+
+    $node = $this->drupalGetNodeByTitle($edit['title']);
+    $node = node_load($node->nid, NULL, TRUE);
+    // Update the status external to our processes.
+    $node->status = 1;
+    $node->title = 'New title';
+    node_save($node);
+
+    $expected = array(
+      'nid' => 1,
+      'vid' => 1,
+      'title' => $node->title,
+      'previous_state' => 'draft',
+      'new_state' => 'published',
+      'status' => 1,
+    );
+    $this->assertTransition($expected);
+
+    // Ensure that multiple saves that do not spawn a new PHP request are
+    // handled correctly to protect against stale static cache.
+    $node = $this->drupalGetNodeByTitle($node->title);
+    $node = node_load($node->nid, NULL, TRUE);
+
+    $node->status = 1;
+    $node->title = 'Newer title';
+    node_save($node);
+    $expected = array(
+      'nid' => 1,
+      'vid' => 1,
+      'title' => $node->title,
+      'previous_state' => 'published',
+      'new_state' => 'published',
+      'status' => 1,
+    );
+    $this->assertTransition($expected);
+  }
+
+}
diff --git a/tests/workbench_moderation_node_access_test/workbench_moderation_node_access_test.info b/tests/workbench_moderation_node_access_test/workbench_moderation_node_access_test.info
new file mode 100644
index 0000000..a4e6f2f
--- /dev/null
+++ b/tests/workbench_moderation_node_access_test/workbench_moderation_node_access_test.info
@@ -0,0 +1,5 @@
+name = Workbench Moderation Node Access Test
+description = Node Access test module for Workbench Moderation.
+package = Workbench
+core = 7.x
+hidden = TRUE
diff --git a/tests/workbench_moderation_node_access_test/workbench_moderation_node_access_test.module b/tests/workbench_moderation_node_access_test/workbench_moderation_node_access_test.module
new file mode 100644
index 0000000..aa77759
--- /dev/null
+++ b/tests/workbench_moderation_node_access_test/workbench_moderation_node_access_test.module
@@ -0,0 +1,35 @@
+<?php
+
+/**
+ * @file
+ * Test module for node access control on moderated nodes.
+ */
+
+/**
+ * Implements hook_node_grants().
+ */
+function workbench_moderation_node_access_test_node_grants($account, $op) {
+  // Give everyone full grants so we don't break other node tests.
+  return array(
+    'test_wm_realm' => array(1),
+  );
+}
+
+/**
+ * Implements hook_node_access_records().
+ */
+function workbench_moderation_node_access_test_node_access_records($node) {
+  // Return nothing for unpublished nodes.
+  if (!$node->status) {
+    return array();
+  }
+  $grants[] = array(
+    'realm' => 'test_wm_realm',
+    'gid' => 1,
+    'grant_view' => 1,
+    'grant_update' => 0,
+    'grant_delete' => 0,
+    'priority' => 0,
+  );
+  return $grants;
+}
diff --git a/tests/workbench_moderation_test.info b/tests/workbench_moderation_test.info
deleted file mode 100644
index 0f9e013..0000000
--- a/tests/workbench_moderation_test.info
+++ /dev/null
@@ -1,5 +0,0 @@
-name = Workbench Moderation Test
-description = Test module for Workbench Moderation.
-package = Workbench
-core = 7.x
-hidden = TRUE
diff --git a/tests/workbench_moderation_test.module b/tests/workbench_moderation_test.module
deleted file mode 100644
index f1dc668..0000000
--- a/tests/workbench_moderation_test.module
+++ /dev/null
@@ -1,37 +0,0 @@
-<?php
-
-/**
- * @file
- * Test module for Workbench Moderation.
- */
-
-/**
- * Implements hook_menu().
- */
-function workbench_moderation_test_menu() {
-  return array(
-    'workbench_moderation_test/%node' => array(
-      'title' => 'Publish a node',
-      'page callback' => 'workbench_moderation_test_update_node',
-      'page arguments' => array(1),
-      'access arguments' => array('bypass workbench moderation'),
-    ),
-  );
-}
-
-/**
- * Page callback. Publishes, unpublishes or resaves the given node.
- *
- * @param object $node
- *   The node to publish, unpublish or resave.
- * @param string $action
- *   Optionally the action to take, either 'publish' or 'unpublish'. If omitted
- *   the node will be resaved.
- */
-function workbench_moderation_test_update_node($node, $action = NULL) {
-  if (!empty($action)) {
-    $node->status = $action == 'publish' ? NODE_PUBLISHED : NODE_NOT_PUBLISHED;
-  }
-  node_save($node);
-  return array('#markup' => t('Node status: @status', array('@status' => $node->status ? t('published') : t('unpublished'))));
-}
diff --git a/tests/workbench_moderation_test/workbench_moderation_test.info b/tests/workbench_moderation_test/workbench_moderation_test.info
new file mode 100644
index 0000000..0f9e013
--- /dev/null
+++ b/tests/workbench_moderation_test/workbench_moderation_test.info
@@ -0,0 +1,5 @@
+name = Workbench Moderation Test
+description = Test module for Workbench Moderation.
+package = Workbench
+core = 7.x
+hidden = TRUE
diff --git a/tests/workbench_moderation_test/workbench_moderation_test.module b/tests/workbench_moderation_test/workbench_moderation_test.module
new file mode 100644
index 0000000..db57e61
--- /dev/null
+++ b/tests/workbench_moderation_test/workbench_moderation_test.module
@@ -0,0 +1,81 @@
+<?php
+
+/**
+ * @file
+ * Test module for Workbench Moderation.
+ */
+
+/**
+ * Implements hook_menu().
+ */
+function workbench_moderation_test_menu() {
+  return array(
+    'workbench_moderation_test/%node' => array(
+      'title' => 'Publish a node',
+      'page callback' => 'workbench_moderation_test_update_node',
+      'page arguments' => array(1),
+      'access arguments' => array('bypass workbench moderation'),
+    ),
+  );
+}
+
+/**
+ * Page callback. Publishes, unpublishes or resaves the given node.
+ *
+ * @param object $node
+ *   The node to publish, unpublish or resave.
+ * @param string $action
+ *   Optionally the action to take, either 'publish' or 'unpublish'. If omitted
+ *   the node will be resaved.
+ */
+function workbench_moderation_test_update_node($node, $action = NULL) {
+  if (!empty($action)) {
+    $node->status = $action == 'publish' ? NODE_PUBLISHED : NODE_NOT_PUBLISHED;
+  }
+  node_save($node);
+  return array('#markup' => t('Node status: @status', array('@status' => $node->status ? t('published') : t('unpublished'))));
+}
+
+/**
+ * Implements hook_workbench_moderation_transition().
+ */
+function workbench_moderation_test_workbench_moderation_transition($node, $previous_state, $new_state) {
+  foreach (array('title', 'nid', 'vid', 'status') as $key) {
+    workbench_moderation_test_set($key, $node->{$key});
+  }
+  workbench_moderation_test_set('previous_state', $previous_state);
+  workbench_moderation_test_set('new_state', $new_state);
+}
+
+/**
+ * Sets values for testing api hooks.
+ *
+ * @param $name
+ *   The name of the value to store.
+ * @param $value
+ *   The value to store.
+ *
+ * @return string | NULL
+ *   The value of a requested named variable, if requested.
+ */
+function workbench_moderation_test_set($name, $value = NULL) {
+  $values = variable_get('workbench_moderation_test', array());
+  if (!is_null($value)) {
+    $values[$name] = $value;
+    variable_set('workbench_moderation_test', $values);
+  }
+  return isset($values[$name]) ? $values[$name] : NULL;
+}
+
+/**
+ * Gets values for testing api hooks.
+ *
+ * @param $name
+ *   The name of the value to store.
+ *
+ * @return string | NULL
+ *   The value of the variable, if set.
+ */
+function workbench_moderation_test_get($name) {
+  return workbench_moderation_test_set($name);
+}
diff --git a/workbench_moderation.info b/workbench_moderation.info
index 3f9939f..a92c565 100644
--- a/workbench_moderation.info
+++ b/workbench_moderation.info
@@ -16,3 +16,8 @@ files[] = workbench_moderation.migrate.inc
 files[] = tests/external_node_update.test
 files[] = tests/workbench_moderation.test
 files[] = tests/workbench_moderation.files.test
+files[] = tests/workbench_moderation.perms.test
+files[] = tests/workbench_moderation.node_access.test
+files[] = tests/workbench_moderation.transition.test
+
+dependencies[] = drafty
diff --git a/workbench_moderation.module b/workbench_moderation.module
index 5a2ff16..98340ef 100644
--- a/workbench_moderation.module
+++ b/workbench_moderation.module
@@ -598,6 +598,7 @@ function _workbench_moderation_access_current_draft($node) {
   }
 
   $state = $node->workbench_moderation;
+
   return (_workbench_moderation_access('view revisions', $node)
       && !empty($state['published'])
       && $state['published']->vid != $state['current']->vid);
@@ -674,65 +675,70 @@ function workbench_moderation_views_default_views() {
   return $return;
 }
 
-
 /**
- * Implements hook_node_presave().
- *
- * Ensure that a node in moderation has the proper publication status.
- * We set $node->status = 0 (unpublished) if this is a new node which has not
- * been marked as published, or if the node has no published revision.
+ * Implements hook_entity_presave().
  */
-function workbench_moderation_node_presave($node) {
-  global $user;
-  if (isset($node->workbench_moderation_state_new)) {
-    // If the new moderation state is published, set the node status to
-    // published.
-    if ($node->workbench_moderation_state_new == workbench_moderation_state_published()) {
-      $node->status = 1;
-    }
-    else {
-      $node->status = 0;
-    }
+function workbench_moderation_entity_presave($entity, $entity_type) {
+  // Note: this only supports nodes at the moment.
+  if ($entity_type != 'node') {
+    return;
+  }
+  // Is the content type under moderation?
+  if (!workbench_moderation_node_type_moderated($entity->type)) {
+    return;
+  }
+
+  // Load our data onto the entity.
+  if (!$entity->is_new) {
+    workbench_moderation_set_node_state($entity);
+  }
+
+  $published = FALSE;
+  // Ensure the entity is marked as published.
+  if (isset($entity->workbench_moderation_state_new) && $entity->workbench_moderation_state_new == workbench_moderation_state_published()) {
+    $published = TRUE;
+    $entity->status = NODE_PUBLISHED;
   }
 
-  // Preserve the changed timestamp of the revision when updating live revision.
-  if (!empty($node->workbench_moderation['updating_live_revision'])) {
-    $node->timestamp = $node->workbench_moderation['my_revision']->stamp;
-    $node->changed = $node->workbench_moderation['my_revision']->stamp;
+  // If we are unpublishing a node, do not force revision.
+  if (!$entity->status) {
+    $entity->is_draft_revision = FALSE;
   }
+
+  // Inform drafty if this is a new revision, if it has not already been
+  // marked as such.
+  elseif (!isset($entity->is_draft_revision) && !$published) {
+    $entity->is_draft_revision = TRUE;
+  }
+
 }
 
 /**
  * Implements hook_node_insert().
  *
- * Wrapper call to the update hook.
+ * Store moderation data for this node.
  */
 function workbench_moderation_node_insert($node) {
-  workbench_moderation_node_update($node);
+  workbench_moderation_save($node);
 }
 
 /**
- * Implements hook_node_update().
- *
- * Handles the submit of the node form moderation information
+ * Loads moderation information onto a node being saved.
  */
-function workbench_moderation_node_update($node) {
+function workbench_moderation_set_node_state($node) {
   global $user;
 
-  // Don't proceed if moderation is not enabled on this content, or if
-  // we're replacing an already-published revision.
-  if (!workbench_moderation_node_moderated($node) || 
-      !empty($node->workbench_moderation['updating_live_revision'])) {
-    return;
-  }
-
   // Set moderation state values.
   if (!isset($node->workbench_moderation_state_current)) {
     $node->workbench_moderation_state_current = !empty($node->original->workbench_moderation['current']->state) ? $node->original->workbench_moderation['current']->state : workbench_moderation_state_none();
   }
   if (!isset($node->workbench_moderation_state_new)) {
+    // Moderating the default revision. Capture the proper state from drafty.
+    if (!empty($node->default_revision) && $node->status) {
+      $node->workbench_moderation_state_new = workbench_moderation_state_published();
+    }
     // Moving from published to unpublished.
-    if ($node->status == NODE_NOT_PUBLISHED && isset($node->original->status) && $node->original->status == NODE_PUBLISHED) {
+    elseif ($node->status == NODE_NOT_PUBLISHED && isset($node->original->status) && $node->original->status == NODE_PUBLISHED) {
       // @todo Currently we cannot set the state correctly if the default state
       //   is "Published".
       // @see https://www.drupal.org/node/1436260
@@ -766,17 +772,85 @@ function workbench_moderation_node_update($node) {
       'stamp' => $node->changed,
     );
   }
+}
 
-  // Apply moderation changes if this is a new revision or if the moderation
-  // state has changed.
-  if (!empty($node->revision) || $node->workbench_moderation_state_current != $node->workbench_moderation_state_new) {
-    // Update attached fields.
-    field_attach_update('node', $node);
-    // Moderate the node.
-    workbench_moderation_moderate($node, $node->workbench_moderation_state_new);
+/**
+ * Implements hook_node_update().
+ *
+ * Stores the moderation information for this node.
+ */
+function workbench_moderation_node_update($node) {
+  workbench_moderation_save($node);
+}
+
+/**
+ * Saves the moderation history for the node.
+ */
+function workbench_moderation_save($node) {
+  $current_draft = &drupal_static(__FUNCTION__, 0);
+  global $user;
+
+  // Don't proceed if moderation is not enabled on this content or if we
+  // are saving the published version from drafty.
+  if (!workbench_moderation_node_moderated($node)) {
+    return;
+  }
+  // Ensure that we have loaded our data onto the node. This function will
+  // check that the required properties are set for all nodes.
+  workbench_moderation_set_node_state($node);
+
+  // Prepare the state information.
+  $state = $node->workbench_moderation_state_new;
+  $old_revision = $node->workbench_moderation['my_revision'];
+
+  $node->is_current = FALSE;
+  if (empty($node->default_revision)) {
+    $node->is_current = TRUE;
+  }
+
+
+  // Build a history record.
+  $new_revision = (object) array(
+    'from_state' => $old_revision->state,
+    'state' => $state,
+    'nid' => $node->nid,
+    'vid' => $node->vid,
+    'uid' => $user->uid,
+    'is_current' => !empty($node->is_current),
+    'published' => ($state == workbench_moderation_state_published()),
+    'stamp' => $_SERVER['REQUEST_TIME'],
+  );
+
+  // If this is the new 'current' moderation record, it should be the only one
+  // flagged 'current' in {workbench_moderation_node_history}.
+  if ($new_revision->is_current) {
+    $query = db_update('workbench_moderation_node_history')
+      ->condition('nid', $node->nid)
+      ->fields(array('is_current' => 0))
+      ->execute();
   }
 
-  return;
+  // If this revision is to be published, the new moderation record should be
+  // the only one flagged 'published' in {workbench_moderation_node_history}.
+  // Also applies in the case where we unpublish a live revision.
+  if ($new_revision->published || !$node->status) {
+    $query = db_update('workbench_moderation_node_history')
+      ->condition('nid', $node->nid)
+      ->fields(array('published' => 0))
+      ->execute();
+  }
+
+  // Save the node history record.
+  drupal_write_record('workbench_moderation_node_history', $new_revision);
+
+  // On a moderation loop, inform other modules of the change.
+  if (!empty($node->is_current)) {
+    // Clear the node's cache.
+    entity_get_controller('node')->resetCache(array($node->nid));
+
+    // Notify other modules that the state was changed.
+    module_invoke_all('workbench_moderation_transition', $node, $node->workbench_moderation['my_revision']->state, $node->workbench_moderation_state_new);
+  }
 }
 
 /**
@@ -1169,13 +1243,12 @@ function workbench_moderation_node_data($node) {
   // We'll store moderation state information in an array on the node.
   $node->workbench_moderation = array();
 
-  // Fetch the most recent revision from the {node_revision} table. This is the
+  // Fetch the most current revision from the {node_revision} table. This is the
   // current revision ("head").
   $query = db_select('node_revision', 'r');
   $query->addJoin('LEFT OUTER', 'workbench_moderation_node_history', 'm', 'r.vid = %alias.vid');
   $query->condition('r.nid', $node->nid)
-    ->orderBy('r.vid', 'DESC')
-    ->orderBy('m.hid', 'DESC')
+    ->condition('m.is_current', 1)
     ->fields('m')
     ->fields('r', array('title', 'timestamp'));
   $current = $query->execute()->fetchObject();
@@ -1624,154 +1697,23 @@ function workbench_moderation_states_next($current_state, $account = NULL, $node
  *   The new moderation state requested.
  */
 function workbench_moderation_moderate($node, $state) {
-  global $user;
-
-  $old_revision = $node->workbench_moderation['my_revision'];
-
-  // Get the number of revisions for this node with vids greater than $node->vid
-  $vid_count = db_select('node_revision', 'r')
-    ->condition('r.nid', $node->nid)
-    ->condition('r.vid', $node->vid, '>')
-    ->countQuery()->execute()->fetchField();
-  // If the number of greater vids is 0, then this is the most current revision
-  $current = ($vid_count == 0);
-
-  // Build a history record.
-  $new_revision = (object) array(
-    'from_state' => $old_revision->state,
-    'state' => $state,
-    'nid' => $node->nid,
-    'vid' => $node->vid,
-    'uid' => $user->uid,
-    'is_current' => $current,
-    'published' => ($state == workbench_moderation_state_published()),
-    'stamp' => $_SERVER['REQUEST_TIME'],
-  );
-
-  // If this is the new 'current' moderation record, it should be the only one
-  // flagged 'current' in {workbench_moderation_node_history}.
-  if ($new_revision->is_current) {
-    $query = db_update('workbench_moderation_node_history')
-      ->condition('nid', $node->nid)
-      ->fields(array('is_current' => 0))
-      ->execute();
-  }
-
-  // If this revision is to be published, the new moderation record should be
-  // the only one flagged 'published' in both
-  // {workbench_moderation_node_history} AND {node_revision}
-  if ($new_revision->published) {
-    $query = db_update('workbench_moderation_node_history')
-      ->condition('nid', $node->nid)
-      ->fields(array('published' => 0))
-      ->execute();
-    $query = db_update('node_revision')
-      ->condition('nid', $node->nid)
-      ->fields(array('status' => 0))
-      ->execute();
-  }
+  // Set the current and new moderation state value.
+  $node->original = $node;
 
-  // Save the node history record.
-  drupal_write_record('workbench_moderation_node_history', $new_revision);
-
-  // Update the node's content_moderation information so that we can publish it
-  // if necessary.
-  $node->workbench_moderation['my_revision'] = $new_revision;
-  if ($new_revision->is_current) {
-    $node->workbench_moderation['current'] = $new_revision;
-  }
-  // Handle the published revision.
-  if ($new_revision->published) {
-    // If we're moderating a revision to the published state, mark the new
-    // revision as the published revision.
-    $node->workbench_moderation['published'] = $new_revision;
-  }
-  elseif (isset($node->workbench_moderation['published']) && $new_revision->vid == $node->workbench_moderation['published']->vid && $new_revision->from_state == workbench_moderation_state_published()) {
-    // If we're moderating the published revision to a non-published state,
-    // remove the workbench moderation 'published' property.
-    $query = db_update('workbench_moderation_node_history')
-      ->condition('hid', $node->workbench_moderation['published']->hid)
-      ->fields(array('published' => 0))
-      ->execute();
-    unset($node->workbench_moderation['published']);
-    $node->workbench_moderation['current']->unpublishing = TRUE;
-  }
-
-  // If we need to make changes to the currently published node we do this in a
-  // shutdown function to avoid race conditions when running node_save() from
-  // within a node submission. We need to change the published node:
-  // - If we're moderating an unpublished revision and there is an existing
-  //   published revision, make sure that the published revision is live.
-  // - If we are moving to unpublished state we should make sure the published
-  //   revision is the 'current' revision.
-  if (!empty($node->workbench_moderation['published']) || !empty($node->workbench_moderation['current']->unpublishing)) {
-    // Clone the node to make sure our data arrives intact in the shutdown
-    // function. It might still be altered before the shutdown is reached.
-    drupal_register_shutdown_function('workbench_moderation_store', clone $node);
-  }
-  else {
-    entity_get_controller('node')->resetCache(array($node->nid));
+  // Ensure that published nodes are flagged properly.
+  if ($state == workbench_moderation_state_published()) {
+    $node->status = NODE_PUBLISHED;
   }
-
-  // Notify other modules that the state was changed.
-  module_invoke_all('workbench_moderation_transition', $node, $old_revision->state, $state);
-}
-
-/**
- * Shutdown callback for saving a node revision.
- *
- * This function is called by drupal_register_shutdown_function().
- * The purpose is to delay a node_save() call so that a live revision
- * is not called during hook_node_update().
- *
- * Instead, we delay the update until the new revision is saved. This way,
- * we can more safely call the revision and pick up changes to items
- * that are not revisioned (such as menu and path assignments).
- *
- * @see workbench_moderation_moderate()
- *
- * @param $node
- *   The node being saved.
- */
-function workbench_moderation_store($node) {
-  if (!isset($node->nid)) {
-    watchdog('Workbench moderation', 'Failed to save node revision: node not passed to shutdown function.', array(), WATCHDOG_NOTICE);
-    return;
+  // If the published version is being unpublished, account for that.
+  elseif (isset($node->workbench_moderation['published']) && $node->workbench_moderation['published']->vid == $node->vid) {
+    $node->status = NODE_NOT_PUBLISHED;
   }
-  watchdog('Workbench moderation', 'Saved node revision: %node as live version for node %live.', array('%node' => $node->vid, '%live' => $node->nid), WATCHDOG_NOTICE, l($node->title, 'node/' . $node->nid));
 
-  // If we are saving a published node, work from the live revision, otherwise
-  // make sure that the entry in the {node} table points to the current
-  // revision.
-  if (empty($node->workbench_moderation['current']->unpublishing)) {
-    $live_revision = workbench_moderation_node_live_load($node);
-    $live_revision->status = 1;
-  }
-  else {
-    $live_revision = workbench_moderation_node_current_load($node);
-    $live_revision->status = 0;
-  }
-  // Don't create a new revision.
-  $live_revision->revision = 0;
-  // Prevent another moderation record from being written.
-  $live_revision->workbench_moderation['updating_live_revision'] = TRUE;
-
-  // Reset flag from taxonomy_field_update() so that {taxonomy_index} values aren't written twice.
-  $taxonomy_index_flag = &drupal_static('taxonomy_field_update', array());
-  unset($taxonomy_index_flag[$node->nid]);
-
-  // Ensure we do not have field translations belonging to a draft revision in
-  // the field data tables.
-  $empty_values = array_fill_keys(array_keys(language_list()), array());
-  foreach (field_info_instances('node', $live_revision->type) as $field_name => $instance) {
-    $field = field_info_field($field_name);
-    if (!empty($live_revision->{$field_name}) && field_is_translatable('node', $field)) {
-      $live_revision->{$field_name} += $empty_values;
-    }
-  }
+  // Set the state property for saving.
+  $node->workbench_moderation_state_new = $state;
 
-  // Save the node.
-  node_save($live_revision);
+  // Save the node and let drafty handle it.
+  node_save($node);
 }
 
 /**
