3 files changed, 41 insertions, 5 deletions

diff --git a/app/models/node.rb b/app/models/node.rb
index 9eb0fe4..f15c908 100644
--- a/app/models/node.rb
+++ b/app/models/node.rb
@@ -56,6 +56,11 @@ class Node < ApplicationRecord
   def lock_for_editing! current_user
     if self.lock_owner.nil? || self.lock_owner == current_user
       lock_for! current_user
+      if self.draft
+        self.draft.user = current_user if self.draft.user.nil?
+        self.draft.editor = current_user
+        self.draft.save!
+      end
       self
     else
       raise(
@@ -95,15 +100,18 @@ class Node < ApplicationRecord
     return unless self.autosave
 
     if self.draft
+      preserved_published_at = self.draft.published_at
       self.draft.clone_attributes_from self.autosave
+      self.draft.published_at = preserved_published_at
       self.draft.user_id = self.autosave.user_id if self.autosave.user_id
       self.draft.editor  = current_user
       self.draft.save!
     else
-      empty_page        = self.pages.create!
-      empty_page.user   = self.autosave.user_id ? self.autosave.user : (self.head ? self.head.user : current_user)
-      empty_page.editor = current_user
+      empty_page = self.pages.create!
       empty_page.clone_attributes_from self.autosave
+      empty_page.user         = self.autosave.user_id ? self.autosave.user : (self.head ? self.head.user : current_user)
+      empty_page.editor       = current_user
+      empty_page.published_at = self.head.published_at if self.head
       empty_page.save!
       self.draft = empty_page
       self.save!
@@ -150,6 +158,28 @@ class Node < ApplicationRecord
     self.draft.reload
   end
 
+  # Discards exactly the topmost non-empty layer -- autosave if present,
+  # else draft -- and reveals whatever's beneath it. Releases the lock
+  # only once nothing is left to protect (no draft survives); leaves it
+  # alone whenever a draft remains, since #edit still has real content
+  # open.
+  def revert! current_user
+    assert_locked_by! current_user
+
+    if self.autosave
+      self.autosave.destroy
+      self.autosave_id = nil
+      self.save!
+    elsif self.draft && self.head
+      self.draft.destroy
+      self.draft_id = nil
+      self.save!
+    end
+
+    self.unlock! unless self.draft
+    self.reload
+  end
+
   def staged_slug=(value)
     if head.blank?
       self.slug = value
diff --git a/test/controllers/nodes_controller_test.rb b/test/controllers/nodes_controller_test.rb
index 030cff0..f3e04c2 100644
--- a/test/controllers/nodes_controller_test.rb
+++ b/test/controllers/nodes_controller_test.rb
@@ -101,7 +101,7 @@ class NodesControllerTest < ActionController::TestCase
     assert_select("#page_body", "World")
 
     node.reload
-    assert_equal 2, node.pages.length
+    assert_equal 1, node.pages.length
     assert_equal "Hello", node.find_or_create_draft( User.first ).title
     assert_equal "World", node.find_or_create_draft( User.first ).body
   end
@@ -291,6 +291,7 @@ class NodesControllerTest < ActionController::TestCase
     get :edit, params: { :id => node.id }
     assert_response :success
 
+    put :update,  params: { :id => node.id, :page => { :title => "updated" } }
     put :publish, params: { :id => node.id }
 
     node.reload
@@ -303,6 +304,7 @@ class NodesControllerTest < ActionController::TestCase
     node  = create_node_with_published_page
     get :edit,    params: { :id => node.id }
 
+    put :update,  params: { :id => node.id, :page => { :title => "updated" } }
     put :publish, params: { :id => node.id }
 
     node.reload
@@ -324,7 +326,8 @@ class NodesControllerTest < ActionController::TestCase
     assert_equal "quentin", node.head.user.login
 
     login_as :aaron
-    get :edit,  params: {:id => node.id }
+    get :edit,   params: { :id => node.id }
+    put :update, params: { :id => node.id, :page => { :title => "updated" } }
 
     node.reload
     assert_equal "quentin", node.head.user.login
diff --git a/test/models/node_test.rb b/test/models/node_test.rb
index 2953f8f..bdf556d 100644
--- a/test/models/node_test.rb
+++ b/test/models/node_test.rb
@@ -339,6 +339,9 @@ class NodeTest < ActiveSupport::TestCase
     assert_equal head_revision, node.head.revision
     assert_nil node.autosave
     assert_equal 2, node.pages.count
+    assert_equal node.head.user, node.draft.user
+    assert_equal @user1, node.draft.editor
+    assert_equal node.head.published_at, node.draft.published_at
   end
 
   test "autosave!, save_draft!, and lock_for_editing! raise LockedByAnotherUser for a second user" do