Fix/821 match error rendering error pages

lib/beacon/plug.ex

@@ -38,11 +38,19 @@ defmodule Beacon.Plug do
   defp put_roll(conn) do
     path_list = conn.path_params["path"]
 
-    with %{private: %{phoenix_live_view: {_, _, %{extra: %{session: %{"beacon_site" => site}}}}}} <- conn,
+    with %{private: %{phoenix_live_view: {_, _, %{extra: %{session: site_session}}}}} <- conn,
+         site <- fetch_session_site(site_session),
          {_, _} <- Beacon.RouterServer.lookup_path(site, path_list, 1) do
       Plug.Conn.put_session(conn, "beacon_variant_roll", Enum.random(1..100))
     else
       _ -> conn
     end
   end
+
+  defp fetch_session_site(site_session) do
+    case site_session do
+      {Beacon.Router, :session, [site, _]} -> site
+      %{"beacon_site" => site} -> site
+    end
+  end
 end

lib/beacon/router.ex

@@ -323,6 +323,9 @@ defmodule Beacon.Router do
       %{phoenix_live_view: {Beacon.Web.PageLive, _, _, %{extra: %{session: %{"beacon_site" => ^site}}}}} ->
         {:ok, {endpoint, host}}
 
+      %{phoenix_live_view: {Beacon.Web.PageLive, _, _, %{extra: %{session: {Beacon.Router, :session, [^site, _sess]}}}}} ->
+        {:ok, {endpoint, host}}
+
       _ ->
         {:error, {endpoint, host}}
     end

lib/beacon/web/controllers/error_html.ex

@@ -8,7 +8,7 @@ defmodule Beacon.Web.ErrorHTML do
 
   @doc false
   def render(<<status_code::binary-size(3), _rest::binary>> = template, %{conn: conn}) do
-    {_, _, %{extra: %{session: %{"beacon_site" => site}}}} = conn.private.phoenix_live_view
+    site = fetch_session_site(conn)
     error_module = Beacon.Loader.fetch_error_page_module(site)
     conn = Plug.Conn.assign(conn, :beacon, Beacon.Web.BeaconAssigns.new(site))
     Beacon.apply_mfa(site, error_module, :render, [conn, String.to_integer(status_code)])
@@ -29,4 +29,23 @@ defmodule Beacon.Web.ErrorHTML do
     Logger.warning("could not find an error page for #{template}, fallbacking to default Phoenix error page")
     Phoenix.Controller.status_message_from_template(template)
   end
+
+  # To be able to render error templates beacon_live_admin
+  # Only works when editing existin sites. Going to admin will fallback
+  # to phoenix error. Any suggestions on how to handle this, will be nice.
+  defp fetch_session_site(%{params: %{"site" => site}})
+    when is_binary(site) and byte_size(site) > 0 do
+    String.to_atom(site)
+  end
+
+  defp fetch_session_site(%{
+    private: %{
+      phoenix_live_view: {_, _, %{extra: %{session: site_session}}}
+    }
+  }) do
+    case site_session do
+      %{"beacon_site" => site} -> site
+      {Beacon.Router, :session, [site, _]} -> site
+    end
+  end
 end

test/beacon/router_test.exs

@@ -67,6 +67,15 @@ defmodule Beacon.RouterTest do
       )
     end
 
+    test "phoenix live_view site map" do
+      config2 = Map.merge(
+        Beacon.Config.fetch!(:my_site),
+        %{host: "host_test", prefix: "/", router: Beacon.BeaconTest.Router}
+      )
+
+      assert Router.reachable?(config2, host: "host.com", prefix: "/host_test")
+    end
+
     test "match existing host" do
       config = config(:host_test)
       assert Router.reachable?(config, host: "host.com", prefix: "/host_test")

test/beacon_web/controllers/error_html_test.exs

@@ -1,15 +1,41 @@
 defmodule Beacon.Web.ErrorHTMLTest do
-  use ExUnit.Case, async: true
+  use Beacon.Web.ConnCase, async: false
+  use Beacon.Test, site: :my_site
 
   alias Beacon.Web.ErrorHTML
 
-  @tag capture_log: true
-  test "invalid status code" do
-    assert ErrorHTML.render("invalid", %{conn: nil}) == "Internal Server Error"
+  describe "render/2 unit test" do
+    @tag capture_log: true
+    test "invalid status code" do
+      assert ErrorHTML.render("invalid", %{conn: nil}) == "Internal Server Error"
+    end
+
+    @tag capture_log: true
+    test "invalid conn" do
+      assert ErrorHTML.render("404.html", %{conn: nil}) == "Not Found"
+    end
   end
 
-  @tag capture_log: true
-  test "invalid conn" do
-    assert ErrorHTML.render("404.html", %{conn: nil}) == "Not Found"
+  describe "render/2 integration with dynamic error pages" do
+    setup %{conn: conn} do
+      beacon_error_page_fixture(
+        status: 404,
+        template: "My Site Not Found Page"
+      )
+
+      conn =
+        conn
+        |> Plug.Conn.assign(:beacon, %{site: :my_site})
+        |> Plug.Conn.put_private(:phoenix_router, Beacon.BeaconTest.Router)
+
+      {:ok, conn: conn}
+    end
+
+    test "missing path", %{conn: conn} do
+      {404, _headers, body} =
+        assert_error_sent(404, fn -> get(conn, "/missing_path") end)
+
+      assert body =~ "My Site Not Found Page"
+    end
   end
 end