2 files changed, 45 insertions, 6 deletions

diff --git a/app/controllers/shared_previews_controller.rb b/app/controllers/shared_previews_controller.rb
index d482466..f6fb45a 100644
--- a/app/controllers/shared_previews_controller.rb
+++ b/app/controllers/shared_previews_controller.rb
@@ -3,13 +3,17 @@ class SharedPreviewsController < ApplicationController
     @page = Page.find_by!(preview_token: params[:token])
     node  = @page.node
 
-    was_published    = @page.published_at.present?
-    superseded       = was_published && node && node.head_id != @page.id
-    currently_public = was_published && node && node.head_id == @page.id && @page.public?
+    if node
+      is_head  = node.head_id  == @page.id
+      is_draft = node.draft_id == @page.id
 
-    if superseded || currently_public
-      redirect_to node_path(node)
-      return
+      currently_public = is_head && @page.public?
+      superseded        = !is_head && !is_draft
+
+      if superseded || currently_public
+        redirect_to @page.public_link
+        return
+      end
     end
 
     response.headers['X-Robots-Tag'] = 'noindex'
diff --git a/test/controllers/shared_previews_controller_test.rb b/test/controllers/shared_previews_controller_test.rb
new file mode 100644
index 0000000..5f2d20d
--- /dev/null
+++ b/test/controllers/shared_previews_controller_test.rb
@@ -0,0 +1,35 @@
+require 'test_helper'
+
+class SharedPreviewsControllerTest < ActionController::TestCase
+  test "renders the preview for a draft that is current but not yet head" do
+    node = Node.root.children.create!(:slug => "shared_preview_draft_test")
+    node.draft.ensure_preview_token!
+
+    get :show, params: { :token => node.draft.preview_token }
+
+    assert_response :success
+  end
+
+  test "renders the preview for a brand-new draft on an already-published node" do
+    node = Node.root.children.create!(:slug => "shared_preview_published_node_test")
+    node.publish_draft!
+    node.find_or_create_draft(User.first)
+    node.draft.ensure_preview_token!
+
+    get :show, params: { :token => node.draft.preview_token }
+
+    assert_response :success
+  end
+
+  test "redirects to the live page once the previewed draft has been published and become head" do
+    node = Node.root.children.create!(:slug => "shared_preview_now_head_test")
+    node.draft.ensure_preview_token!
+    token = node.draft.preview_token
+
+    node.publish_draft!
+
+    get :show, params: { :token => token }
+
+    assert_redirected_to node.head.public_link
+  end
+end